datastore: Update AllocateIDs to take keys.

service/datastore/datastore_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.
 
 // adapted from github.com/golang/appengine/datastore
 
 package datastore
 
 import (
         "bytes"
@@ skipping 13 matching lines @@
 
 func fakeDatastoreFactory(c context.Context, wantTxn bool) RawInterface {
         i := info.Get(c)
         fds := fakeDatastore{
                 aid: i.FullyQualifiedAppID(),
         }
         fds.ns, _ = i.GetNamespace()
         return &fds
 }
 
+var (
+        errFail    = errors.New("Individual element fail")
+        errFailAll = errors.New("Operation fail")
+)
+
 type fakeDatastore struct {
         RawInterface
         aid string
         ns  string
 }
 
 func (f *fakeDatastore) mkKey(elems ...interface{}) *Key {
         return MakeKey(f.aid, f.ns, elems...)
 }
 
+func (f *fakeDatastore) newKey(kind, stringID string, intID int64, parent *Key) *Key {
+        return NewKey(f.aid, f.ns, kind, stringID, intID, parent)
+}
+
+func (f *fakeDatastore) AllocateIDs(keys []*Key, cb NewKeyCB) error {
+        if keys[0].Kind() == "FailAll" {
+                return errFailAll
+        }
+        for i, k := range keys {
+                if k.Kind() == "Fail" {
+                        cb(nil, errFail)
+                } else {
+                        cb(f.newKey(k.Kind(), "", int64(i+1), k.Parent()), nil)
+                }
+        }
+        return nil
+}
+
 func (f *fakeDatastore) Run(fq *FinalizedQuery, cb RawRunCB) error {
         lim, _ := fq.Limit()
 
         cursCB := func() (Cursor, error) {
                 return fakeCursor("CURSOR"), nil
         }
 
         for i := int32(0); i < lim; i++ {
                 if v, ok := fq.eqFilts["$err_single"]; ok {
                         idx := fq.eqFilts["$err_single_idx"][0].Value().(int64)
                         if idx == int64(i) {
                                 return errors.New(v[0].Value().(string))
                         }
                 }
                 k := f.mkKey("Kind", i+1)
                 if i == 10 {
                         k = f.mkKey("Kind", "eleven")
                 }
                 pm := PropertyMap{"Value": {MkProperty(i)}}
                 if err := cb(k, pm, cursCB); err != nil {
                         if err == Stop {
                                 err = nil
                         }
                         return err
                 }
         }
         return nil
 }
 
-var (
-        errFail    = errors.New("Individual element fail")
-        errFailAll = errors.New("Operation fail")
-)
-
-func (f *fakeDatastore) PutMulti(keys []*Key, vals []PropertyMap, cb PutMultiCB) error {
+func (f *fakeDatastore) PutMulti(keys []*Key, vals []PropertyMap, cb NewKeyCB) error {
         if keys[0].Kind() == "FailAll" {
                 return errFailAll
         }
         _, assertExtra := vals[0].GetMeta("assertExtra")
         for i, k := range keys {
                 err := error(nil)
                 if k.Kind() == "Fail" {
                         err = errFail
                 } else {
                         So(vals[i]["Value"], ShouldResemble, []Property{MkProperty(i)})
                         if assertExtra {
                                 So(vals[i]["Extra"], ShouldResemble, []Property{MkProperty("whoa")})
                         }
-                        if k.Incomplete() {
+                        if k.IsIncomplete() {
                                 k = NewKey(k.AppID(), k.Namespace(), k.Kind(), "", int64(i+1), k.Parent())
                         }
                 }
                 cb(k, err)
         }
         return nil
 }
 
 const noSuchEntityID = 0xdead
 
@@ skipping 286 matching lines @@
 
 func TestPopulateKey(t *testing.T) {
         t.Parallel()
 
         Convey("Test PopulateKey", t, func() {
                 k := NewKey("app", "namespace", "kind", "", 1337, nil)
 
                 Convey("Can set the key of a common struct.", func() {
                         var cs CommonStruct
 
-                        PopulateKey(&cs, k)
+                        So(PopulateKey(&cs, k), ShouldBeTrue)
                         So(cs.ID, ShouldEqual, 1337)
                 })
 
                 Convey("Will not set the value of a singleton struct.", func() {
                         var ss SingletonStruct
 
-                        PopulateKey(&ss, k)
+                        So(PopulateKey(&ss, k), ShouldBeFalse)
                         So(ss.id, ShouldEqual, 0)
                 })
 
                 Convey("Will panic when setting the key of a bad struct.", func() {
                         var bs badStruct
 
                         So(func() { PopulateKey(&bs, k) }, ShouldPanic)
                 })
 
                 Convey("Will panic when setting the key of a broken PLS struct.", func() {
                         var broken permaBad
 
                         So(func() { PopulateKey(&broken, k) }, ShouldPanic)
                 })
         })
 }
 
+func TestAllocateIDs(t *testing.T) {
+        t.Parallel()
+
+        Convey("A testing environment", t, func() {
+                c := info.Set(context.Background(), fakeInfo{})
+                c = SetRawFactory(c, fakeDatastoreFactory)
+                ds := Get(c)
+
+                Convey("Testing AllocateIDs", func() {
+                        Convey("Will return nil if no entities are supplied.", func() {
+                                So(ds.AllocateIDs(), ShouldBeNil)
+                        })
+
+                        Convey("single struct", func() {
+                                cs := CommonStruct{Value: 1}
+                                So(ds.AllocateIDs(&cs), ShouldBeNil)
+                                So(cs.ID, ShouldEqual, 1)
+                        })
+
+                        Convey("struct slice", func() {
+                                csSlice := []*CommonStruct{{Value: 1}, {Value: 2}}
+                                So(ds.AllocateIDs(csSlice), ShouldBeNil)
+                                So(csSlice, ShouldResemble, []*CommonStruct{{ID: 1, Value: 1}, {ID: 2, Value: 2}})
+                        })
+
+                        Convey("single key will fail", func() {
+                                singleKey := ds.MakeKey("FooParent", "BarParent", "Foo", "Bar")
+                                So(func() { ds.AllocateIDs(singleKey) }, ShouldPanicLike,
+                                        "invalid input type (*datastore.Key): not a PLS, pointer-to-struct, or slice thereof")
+                        })
+
+                        Convey("key slice", func() {
+                                k0 := ds.MakeKey("Foo", "Bar")
+                                k1 := ds.MakeKey("Baz", "Qux")
+                                keySlice := []*Key{k0, k1}
+                                So(ds.AllocateIDs(keySlice), ShouldBeNil)
+                                So(keySlice[0].Equal(ds.MakeKey("Foo", 1)), ShouldBeTrue)
+                                So(keySlice[1].Equal(ds.MakeKey("Baz", 2)), ShouldBeTrue)
+
+                                // The original keys should not have changed.
+                                So(k0.Equal(ds.MakeKey("Foo", "Bar")), ShouldBeTrue)
+                                So(k1.Equal(ds.MakeKey("Baz", "Qux")), ShouldBeTrue)
+                        })
+
+                        Convey("fail all key slice", func() {
+                                keySlice := []*Key{ds.MakeKey("FailAll", "oops"), ds.MakeKey("Baz", "Qux")}
+                                So(ds.AllocateIDs(keySlice), ShouldEqual, errFailAll)
+                                So(keySlice[0].StringID(), ShouldEqual, "oops")
+                                So(keySlice[1].StringID(), ShouldEqual, "Qux")
+                        })
+
+                        Convey("fail key slice", func() {
+                                keySlice := []*Key{ds.MakeKey("Fail", "oops"), ds.MakeKey("Baz", "Qux")}
+                                So(ds.AllocateIDs(keySlice), ShouldResemble, errors.MultiError{errFail, nil})
+                                So(keySlice[0].StringID(), ShouldEqual, "oops")
+                                So(keySlice[1].IntID(), ShouldEqual, 2)
+                        })
+
+                        Convey("vararg with errors", func() {
+                                successSlice := []CommonStruct{{Value: 0}, {Value: 1}}
+                                failSlice := []FakePLS{{Kind: "Fail"}, {Value: 3}}
+                                emptySlice := []CommonStruct(nil)
+                                cs0 := CommonStruct{Value: 4}
+                                cs1 := FakePLS{Kind: "Fail", Value: 5}
+                                keySlice := []*Key{ds.MakeKey("Foo", "Bar"), ds.MakeKey("Baz", "Qux")}
+                                fpls := FakePLS{StringID: "ohai", Value: 6}
+
+                                err := ds.AllocateIDs(successSlice, failSlice, emptySlice, &cs0, &cs1, keySlice, &fpls)
+                                So(err, ShouldResemble, errors.MultiError{
+                                        nil, errors.MultiError{errFail, nil}, nil, nil, errFail, nil, nil})
+                                So(successSlice[0].ID, ShouldEqual, 1)
+                                So(successSlice[1].ID, ShouldEqual, 2)
+                                So(failSlice[1].IntID, ShouldEqual, 4)
+                                So(cs0.ID, ShouldEqual, 5)
+                                So(keySlice[0].Equal(ds.MakeKey("Foo", 7)), ShouldBeTrue)
+                                So(keySlice[1].Equal(ds.MakeKey("Baz", 8)), ShouldBeTrue)
+                                So(fpls.IntID, ShouldEqual, 9)
+                        })
+                })
+        })
+}
+
 func TestPut(t *testing.T) {
         t.Parallel()
 
         Convey("A testing environment", t, func() {
                 c := info.Set(context.Background(), fakeInfo{})
                 c = SetRawFactory(c, fakeDatastoreFactory)
                 ds := Get(c)
 
                 Convey("Testing Put", func() {
                         Convey("bad", func() {
                                 Convey("static can't serialize", func() {
                                         bss := []badStruct{{}, {}}
                                         So(func() { ds.Put(bss) }, ShouldPanicLike,
                                                 `field "Compy" has invalid type`)
                                 })
 
                                 Convey("static ptr can't serialize", func() {
                                         bss := []*badStruct{{}, {}}
                                         So(func() { ds.Put(bss) }, ShouldPanicLike,
                                                 `field "Compy" has invalid type: complex64`)
                                 })
 
                                 Convey("static bad type", func() {
                                         So(func() { ds.Put(100) }, ShouldPanicLike,
-                                                "invalid input type (int): not a PLS or pointer-to-struct")
+                                                "invalid input type (int): not a PLS, pointer-to-struct, or slice thereof")
                                 })
 
                                 Convey("static bad type (slice of bad type)", func() {
                                         So(func() { ds.Put([]int{}) }, ShouldPanicLike,
-                                                "invalid input type ([]int): not a PLS or pointer-to-struct")
+                                                "invalid input type ([]int): not a PLS, pointer-to-struct, or slice thereof")
                                 })
 
                                 Convey("dynamic can't serialize", func() {
                                         fplss := []FakePLS{{failSave: true}, {}}
                                         So(ds.Put(fplss), ShouldErrLike, "FakePLS.Save")
                                 })
 
                                 Convey("can't get keys", func() {
                                         fplss := []FakePLS{{failGetMeta: true}, {}}
                                         So(ds.Put(fplss), ShouldErrLike, "unable to extract $kind")
                                 })
 
                                 Convey("get single error for RPC failure", func() {
                                         fplss := []FakePLS{{Kind: "FailAll"}, {}}
                                         So(ds.Put(fplss), ShouldEqual, errFailAll)
                                 })
 
                                 Convey("get multi error for individual failures", func() {
                                         fplss := []FakePLS{{}, {Kind: "Fail"}}
                                         So(ds.Put(fplss), ShouldResemble, errors.MultiError{nil, errFail})
                                 })
 
                                 Convey("get with *Key is an error", func() {
                                         So(func() { ds.Get(&Key{}) }, ShouldPanicLike,
-                                                "invalid input type (*datastore.Key): not user datatype")
+                                                "invalid input type (*datastore.Key): not a PLS, pointer-to-struct, or slice thereof")
                                 })
 
                                 Convey("struct with no $kind is an error", func() {
                                         s := MGSWithNoKind{}
                                         So(ds.Put(&s), ShouldErrLike, "unable to extract $kind")
                                 })
 
                                 Convey("struct with invalid but non-nil key is an error", func() {
                                         type BadParent struct {
                                                 ID     int64 `gae:"$id"`
@@ skipping 389 matching lines @@
                                 })
 
                                 Convey("get with nil is an error", func() {
                                         So(func() { ds.Get(nil) }, ShouldPanicLike,
                                                 "cannot use nil as single argument")
                                 })
 
                                 Convey("get with ptr-to-nonstruct is an error", func() {
                                         val := 100
                                         So(func() { ds.Get(&val) }, ShouldPanicLike,
-                                                "invalid input type (*int): not a PLS or pointer-to-struct")
+                                                "invalid input type (*int): not a PLS, pointer-to-struct, or slice thereof")
                                 })
 
                                 Convey("failure to save metadata is no problem though", func() {
                                         // It just won't save the key
                                         cs := &FakePLS{IntID: 10, failSetMeta: true}
                                         So(ds.Get(cs), ShouldBeNil)
                                 })
 
                                 Convey("vararg with errors", func() {
                                         successSlice := []CommonStruct{{ID: 1}, {ID: 2}}
@@ skipping 409 matching lines @@
                 data, ok := d.data[k.String()]
                 if ok {
                         cb(data, nil)
                 } else {
                         cb(nil, ErrNoSuchEntity)
                 }
         }
         return nil
 }
 
-func (d *fixedDataDatastore) PutMulti(keys []*Key, vals []PropertyMap, cb PutMultiCB) error {
+func (d *fixedDataDatastore) PutMulti(keys []*Key, vals []PropertyMap, cb NewKeyCB) error {
         if d.data == nil {
                 d.data = make(map[string]PropertyMap, len(keys))
         }
         for i, k := range keys {
-                if k.Incomplete() {
+                if k.IsIncomplete() {
                         panic("key is incomplete, don't do that.")
                 }
                 d.data[k.String()], _ = vals[i].Save(false)
                 cb(k, nil)
         }
         return nil
 }
 
 func TestSchemaChange(t *testing.T) {
         t.Parallel()
@@ skipping 480 matching lines @@
                         if err != nil {
                                 panic(fmt.Errorf("failed to find absolute path for `%s`", sameLevelDir))
                         }
 
                         ids, err := FindAndParseIndexYAML(abs)
                         So(err, ShouldBeNil)
                         So(ids[1].Kind, ShouldEqual, "Test Foo")
                 })
         })
 }