Interface update, per-method Contexts.

impl/prod/everything_test.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 // +build appengine
 
 package prod
 
 import (
         "testing"
         "time"
 
         "github.com/luci/gae/service/blobstore"
-        "github.com/luci/gae/service/datastore"
+        ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/info"
-        "github.com/luci/gae/service/memcache"
+        mc "github.com/luci/gae/service/memcache"
+
         "github.com/luci/luci-go/common/logging"
-        . "github.com/smartystreets/goconvey/convey"
+
         "golang.org/x/net/context"
         "google.golang.org/appengine/aetest"
+
+        . "github.com/smartystreets/goconvey/convey"
 )
 
 var (
-        mp   = datastore.MkProperty
-        mpNI = datastore.MkPropertyNI
+        mp   = ds.MkProperty
+        mpNI = ds.MkPropertyNI
 )
 
 type TestStruct struct {
         ID int64 `gae:"$id"`
 
         ValueI  []int64
         ValueB  []bool
         ValueS  []string
         ValueF  []float64
         ValueBS [][]byte // "ByteString"
-        ValueK  []*datastore.Key
+        ValueK  []*ds.Key
         ValueBK []blobstore.Key
-        ValueGP []datastore.GeoPoint
+        ValueGP []ds.GeoPoint
 }
 
 func TestBasicDatastore(t *testing.T) {
         t.Parallel()
 
         Convey("basic", t, func() {
                 inst, err := aetest.NewInstance(&aetest.Options{
                         StronglyConsistentDatastore: true,
                 })
                 So(err, ShouldBeNil)
                 defer inst.Close()
 
                 req, err := inst.NewRequest("GET", "/", nil)
                 So(err, ShouldBeNil)
 
                 ctx := Use(context.Background(), req)
-                ds := datastore.Get(ctx)
-                mc := memcache.Get(ctx)
-                inf := info.Get(ctx)
 
                 Convey("logging allows you to tweak the level", func() {
                         // You have to visually confirm that this actually happens in the stdout
                         // of the test... yeah I know.
                         logging.Debugf(ctx, "SHOULD NOT SEE")
                         logging.Infof(ctx, "SHOULD SEE")
 
                         ctx = logging.SetLevel(ctx, logging.Debug)
                         logging.Debugf(ctx, "SHOULD SEE")
                         logging.Infof(ctx, "SHOULD SEE (2)")
                 })
 
                 Convey("Can probe/change Namespace", func() {
-                        ns, has := inf.GetNamespace()
-                        So(ns, ShouldEqual, "")
-                        So(has, ShouldBeFalse)
+                        So(info.GetNamespace(ctx), ShouldEqual, "")
 
-                        ctx, err = inf.Namespace("wat")
+                        ctx, err = info.Namespace(ctx, "wat")
                         So(err, ShouldBeNil)
-                        inf = info.Get(ctx)
 
-                        ns, has = inf.GetNamespace()
-                        So(ns, ShouldEqual, "wat")
-                        So(has, ShouldBeTrue)
-
-                        ds = datastore.Get(ctx)
-                        So(ds.MakeKey("Hello", "world").Namespace(), ShouldEqual, "wat")
+                        So(info.GetNamespace(ctx), ShouldEqual, "wat")
+                        So(ds.MakeKey(ctx, "Hello", "world").Namespace(), ShouldEqual, "wat")
                 })
 
                 Convey("Can get non-transactional context", func() {
-                        ctx, err := inf.Namespace("foo")
+                        ctx, err := info.Namespace(ctx, "foo")
                         So(err, ShouldBeNil)
-                        ds = datastore.Get(ctx)
-                        inf = info.Get(ctx)
 
-                        ds.RunInTransaction(func(ctx context.Context) error {
-                                So(ds.MakeKey("Foo", "bar").Namespace(), ShouldEqual, "foo")
+                        So(ds.CurrentTransaction(ctx), ShouldBeNil)
 
-                                So(ds.Put(&TestStruct{ValueI: []int64{100}}), ShouldBeNil)
+                        ds.RunInTransaction(ctx, func(ctx context.Context) error {
+                                So(ds.CurrentTransaction(ctx), ShouldNotBeNil)
+                                So(ds.MakeKey(ctx, "Foo", "bar").Namespace(), ShouldEqual, "foo")
 
-                                err = datastore.GetNoTxn(ctx).RunInTransaction(func(ctx context.Context) error {
-                                        ds = datastore.Get(ctx)
-                                        So(ds.MakeKey("Foo", "bar").Namespace(), ShouldEqual, "foo")
-                                        So(ds.Put(&TestStruct{ValueI: []int64{100}}), ShouldBeNil)
+                                So(ds.Put(ctx, &TestStruct{ValueI: []int64{100}}), ShouldBeNil)
+
+                                noTxnCtx := ds.WithoutTransaction(ctx)
+                                So(ds.CurrentTransaction(noTxnCtx), ShouldBeNil)
+
+                                err = ds.RunInTransaction(noTxnCtx, func(ctx context.Context) error {
+                                        So(ds.CurrentTransaction(ctx), ShouldNotBeNil)
+                                        So(ds.MakeKey(ctx, "Foo", "bar").Namespace(), ShouldEqual, "foo")
+                                        So(ds.Put(ctx, &TestStruct{ValueI: []int64{100}}), ShouldBeNil)
                                         return nil
                                 }, nil)
                                 So(err, ShouldBeNil)
 
                                 return nil
                         }, nil)
                 })
 
                 Convey("Can Put/Get", func() {
                         orig := TestStruct{
                                 ValueI: []int64{1, 7, 946688461000000, 996688461000000},
                                 ValueB: []bool{true, false},
                                 ValueS: []string{"hello", "world"},
                                 ValueF: []float64{1.0, 7.0, 946688461000000.0, 996688461000000.0},
                                 ValueBS: [][]byte{
                                         []byte("allo"),
                                         []byte("hello"),
                                         []byte("world"),
                                         []byte("zurple"),
                                 },
-                                ValueK: []*datastore.Key{
-                                        ds.NewKey("Something", "Cool", 0, nil),
-                                        ds.NewKey("Something", "", 1, nil),
-                                        ds.NewKey("Something", "Recursive", 0,
-                                                ds.NewKey("Parent", "", 2, nil)),
+                                ValueK: []*ds.Key{
+                                        ds.NewKey(ctx, "Something", "Cool", 0, nil),
+                                        ds.NewKey(ctx, "Something", "", 1, nil),
+                                        ds.NewKey(ctx, "Something", "Recursive", 0,
+                                                ds.NewKey(ctx, "Parent", "", 2, nil)),
                                 },
                                 ValueBK: []blobstore.Key{"bellow", "hello"},
-                                ValueGP: []datastore.GeoPoint{
+                                ValueGP: []ds.GeoPoint{
                                         {Lat: 120.7, Lng: 95.5},
                                 },
                         }
-                        So(ds.Put(&orig), ShouldBeNil)
+                        So(ds.Put(ctx, &orig), ShouldBeNil)
 
                         ret := TestStruct{ID: orig.ID}
-                        So(ds.Get(&ret), ShouldBeNil)
+                        So(ds.Get(ctx, &ret), ShouldBeNil)
                         So(ret, ShouldResemble, orig)
 
                         // can't be sure the indexes have caught up... so sleep
                         time.Sleep(time.Second)
 
                         Convey("Can query", func() {
-                                q := datastore.NewQuery("TestStruct")
-                                ds.Run(q, func(ts *TestStruct) {
+                                q := ds.NewQuery("TestStruct")
+                                ds.Run(ctx, q, func(ts *TestStruct) {
                                         So(*ts, ShouldResemble, orig)
                                 })
-                                count, err := ds.Count(q)
+                                count, err := ds.Count(ctx, q)
                                 So(err, ShouldBeNil)
                                 So(count, ShouldEqual, 1)
                         })
 
                         Convey("Can project", func() {
-                                q := datastore.NewQuery("TestStruct").Project("ValueS")
-                                rslts := []datastore.PropertyMap{}
-                                So(ds.GetAll(q, &rslts), ShouldBeNil)
-                                So(rslts, ShouldResemble, []datastore.PropertyMap{
+                                q := ds.NewQuery("TestStruct").Project("ValueS")
+                                rslts := []ds.PropertyMap{}
+                                So(ds.GetAll(ctx, q, &rslts), ShouldBeNil)
+                                So(rslts, ShouldResemble, []ds.PropertyMap{
                                         {
-                                                "$key":   {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":   {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueS": {mp("hello")},
                                         },
                                         {
-                                                "$key":   {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":   {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueS": {mp("world")},
                                         },
                                 })
 
-                                q = datastore.NewQuery("TestStruct").Project("ValueBS")
-                                rslts = []datastore.PropertyMap{}
-                                So(ds.GetAll(q, &rslts), ShouldBeNil)
-                                So(rslts, ShouldResemble, []datastore.PropertyMap{
+                                q = ds.NewQuery("TestStruct").Project("ValueBS")
+                                rslts = []ds.PropertyMap{}
+                                So(ds.GetAll(ctx, q, &rslts), ShouldBeNil)
+                                So(rslts, ShouldResemble, []ds.PropertyMap{
                                         {
-                                                "$key":    {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":    {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueBS": {mp("allo")},
                                         },
                                         {
-                                                "$key":    {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":    {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueBS": {mp("hello")},
                                         },
                                         {
-                                                "$key":    {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":    {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueBS": {mp("world")},
                                         },
                                         {
-                                                "$key":    {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":    {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueBS": {mp("zurple")},
                                         },
                                 })
 
-                                count, err := ds.Count(q)
+                                count, err := ds.Count(ctx, q)
                                 So(err, ShouldBeNil)
                                 So(count, ShouldEqual, 4)
 
-                                q = datastore.NewQuery("TestStruct").Lte("ValueI", 7).Project("ValueS").Distinct(true)
-                                rslts = []datastore.PropertyMap{}
-                                So(ds.GetAll(q, &rslts), ShouldBeNil)
-                                So(rslts, ShouldResemble, []datastore.PropertyMap{
+                                q = ds.NewQuery("TestStruct").Lte("ValueI", 7).Project("ValueS").Distinct(true)
+                                rslts = []ds.PropertyMap{}
+                                So(ds.GetAll(ctx, q, &rslts), ShouldBeNil)
+                                So(rslts, ShouldResemble, []ds.PropertyMap{
                                         {
-                                                "$key":   {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":   {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueI": {mp(1)},
                                                 "ValueS": {mp("hello")},
                                         },
                                         {
-                                                "$key":   {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":   {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueI": {mp(1)},
                                                 "ValueS": {mp("world")},
                                         },
                                         {
-                                                "$key":   {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":   {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueI": {mp(7)},
                                                 "ValueS": {mp("hello")},
                                         },
                                         {
-                                                "$key":   {mpNI(ds.KeyForObj(&orig))},
+                                                "$key":   {mpNI(ds.KeyForObj(ctx, &orig))},
                                                 "ValueI": {mp(7)},
                                                 "ValueS": {mp("world")},
                                         },
                                 })
 
-                                count, err = ds.Count(q)
+                                count, err = ds.Count(ctx, q)
                                 So(err, ShouldBeNil)
                                 So(count, ShouldEqual, 4)
                         })
                 })
 
                 Convey("Can Put/Get (time)", func() {
                         // time comparisons in Go are wonky, so this is pulled out
-                        pm := datastore.PropertyMap{
-                                "$key": {mpNI(ds.NewKey("Something", "value", 0, nil))},
+                        pm := ds.PropertyMap{
+                                "$key": {mpNI(ds.NewKey(ctx, "Something", "value", 0, nil))},
                                 "Time": {
                                         mp(time.Date(1938, time.January, 1, 1, 1, 1, 1, time.UTC)),
                                         mp(time.Time{}),
                                 },
                         }
-                        So(ds.Put(&pm), ShouldBeNil)
+                        So(ds.Put(ctx, &pm), ShouldBeNil)
 
-                        rslt := datastore.PropertyMap{}
-                        rslt.SetMeta("key", ds.KeyForObj(pm))
-                        So(ds.Get(&rslt), ShouldBeNil)
+                        rslt := ds.PropertyMap{}
+                        rslt.SetMeta("key", ds.KeyForObj(ctx, pm))
+                        So(ds.Get(ctx, &rslt), ShouldBeNil)
 
                         So(pm["Time"][0].Value(), ShouldResemble, rslt["Time"][0].Value())
 
-                        q := datastore.NewQuery("Something").Project("Time")
-                        all := []datastore.PropertyMap{}
-                        So(ds.GetAll(q, &all), ShouldBeNil)
+                        q := ds.NewQuery("Something").Project("Time")
+                        all := []ds.PropertyMap{}
+                        So(ds.GetAll(ctx, q, &all), ShouldBeNil)
                         So(len(all), ShouldEqual, 2)
                         prop := all[0]["Time"][0]
-                        So(prop.Type(), ShouldEqual, datastore.PTInt)
+                        So(prop.Type(), ShouldEqual, ds.PTInt)
 
-                        tval, err := prop.Project(datastore.PTTime)
+                        tval, err := prop.Project(ds.PTTime)
                         So(err, ShouldBeNil)
                         So(tval, ShouldResemble, time.Time{}.UTC())
 
-                        tval, err = all[1]["Time"][0].Project(datastore.PTTime)
+                        tval, err = all[1]["Time"][0].Project(ds.PTTime)
                         So(err, ShouldBeNil)
                         So(tval, ShouldResemble, pm["Time"][0].Value())
 
-                        ent := datastore.PropertyMap{
-                                "$key": {mpNI(ds.MakeKey("Something", "value"))},
+                        ent := ds.PropertyMap{
+                                "$key": {mpNI(ds.MakeKey(ctx, "Something", "value"))},
                         }
-                        So(ds.Get(&ent), ShouldBeNil)
+                        So(ds.Get(ctx, &ent), ShouldBeNil)
                         So(ent["Time"], ShouldResemble, pm["Time"])
                 })
 
                 Convey("memcache: Set (nil) is the same as Set ([]byte{})", func() {
-                        So(mc.Set(mc.NewItem("bob")), ShouldBeNil) // normally would panic because Value is nil
+                        So(mc.Set(ctx, mc.NewItem(ctx, "bob")), ShouldBeNil) // normally would panic because Value is nil
 
-                        bob, err := mc.Get("bob")
+                        bob, err := mc.GetKey(ctx, "bob")
                         So(err, ShouldBeNil)
                         So(bob.Value(), ShouldResemble, []byte{})
                 })
         })
 }