datastore: Update AllocateIDs to take keys.

impl/memory/datastore_test.go

Index: impl/memory/datastore_test.go
diff --git a/impl/memory/datastore_test.go b/impl/memory/datastore_test.go
index 3b9f91e212dfc2f75d09512216539fc0785e295e..2480aeef10ff52166d6933821a3437bf0a03ec0e 100644
--- a/impl/memory/datastore_test.go
+++ b/impl/memory/datastore_test.go
@@ -162,14 +162,26 @@ func TestDatastoreSingleReadWriter(t *testing.T) {
 			})
 
 			Convey("allocating ids prevents their use", func() {
-				start, err := ds.AllocateIDs(ds.MakeKey("Foo", 0), 100)
+				keys := ds.NewIncompleteKeys(100, "Foo", nil)
+				keys, err := ds.AllocateIDs(keys...)
 				So(err, ShouldBeNil)
-				So(start, ShouldEqual, 2)
+				So(len(keys), ShouldEqual, 100)
 
+				// Assert that none of our keys share the same ID.
+				ids := make(map[int64]struct{})
+				for _, k := range keys {
+					ids[k.IntID()] = struct{}{}
+				}
+				So(len(ids), ShouldEqual, len(keys))
+
+				// Put a new object and ensure that it is allocated an unused ID.
 				f := &Foo{Val: 10}
 				So(ds.Put(f), ShouldBeNil)
 				k := ds.KeyForObj(f)
 				So(k.String(), ShouldEqual, "dev~app::/Foo,102")
+
+				_, ok := ids[k.IntID()]
+				So(ok, ShouldBeFalse)
 			})
 		})