Add missing Count api

impl/memory/datastore_query_execution_test.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
         "fmt"
         "testing"
         "time"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/info"
         . "github.com/luci/luci-go/common/testing/assertions"
         . "github.com/smartystreets/goconvey/convey"
         "golang.org/x/net/context"
 )
 
 type qExpect struct {
         q     *ds.Query
         inTxn bool
 
-        get  []ds.PropertyMap
-        keys []*ds.Key
+        get   []ds.PropertyMap
+        keys  []*ds.Key
+        count int
 }
 
 type qExStage struct {
         addIdxs []*ds.IndexDefinition
         putEnts []ds.PropertyMap
         delEnts []*ds.Key
 
         expect []qExpect
 
         extraFns []func(context.Context)
@@ skipping 117 matching lines @@
                                 },
 
                                 {q: nq("Kind").Gt("Val", 2).Lte("Val", 5), get: []ds.PropertyMap{
                                         stage1Data[0], stage1Data[3],
                                 }},
 
                                 {q: nq("Kind").Gt("Val", 2).Lte("Val", 5).Order("-Val"), get: []ds.PropertyMap{
                                         stage1Data[3], stage1Data[0],
                                 }},
 
-                                {q: nq("").Gt("__key__", key("Kind", 2)), get: []ds.PropertyMap{
-                                        // TODO(riannucci): determine if the real datastore shows metadata
-                                        // during kindless queries. The documentation seems to imply so, but
-                                        // I'd like to be sure.
-                                        pmap("$key", key("Kind", 2, "__entity_group__", 1), Next,
-                                                "__version__", 1),
-                                        stage1Data[2],
-                                        stage1Data[4],
-                                        // this is 5 because the value is retrieved from HEAD and not from
-                                        // the index snapshot!
-                                        pmap("$key", key("Kind", 3, "__entity_group__", 1), Next,
-                                                "__version__", 5),
-                                        stage1Data[3],
-                                        pmap("$key", key("Kind", 6, "__entity_group__", 1), Next,
-                                                "__version__", 1),
-                                        pmap("$key", key("Unique", 1, "__entity_group__", 1), Next,
-                                                "__version__", 2),
-                                }},
+                                {q: nq("").Gt("__key__", key("Kind", 2)),
+                                        // count counts from the index with KeysOnly and so counts the deleted
+                                        // entity Unique/1.
+                                        count: 8,
+                                        get: []ds.PropertyMap{
+                                                // TODO(riannucci): determine if the real datastore shows metadata
+                                                // during kindless queries. The documentation seems to imply so, but
+                                                // I'd like to be sure.
+                                                pmap("$key", key("Kind", 2, "__entity_group__", 1), Next,
+                                                        "__version__", 1),
+                                                stage1Data[2],
+                                                stage1Data[4],
+                                                // this is 5 because the value is retrieved from HEAD and not from
+                                                // the index snapshot!
+                                                pmap("$key", key("Kind", 3, "__entity_group__", 1), Next,
+                                                        "__version__", 5),
+                                                stage1Data[3],
+                                                pmap("$key", key("Kind", 6, "__entity_group__", 1), Next,
+                                                        "__version__", 1),
+                                                pmap("$key", key("Unique", 1, "__entity_group__", 1), Next,
+                                                        "__version__", 2),
+                                        }},
 
                                 {q: (nq("Kind").
                                         Gt("Val", 2).Eq("Extra", "waffle").
                                         Order("-Val").
                                         Ancestor(key("Kind", 3))),
                                         get: []ds.PropertyMap{
                                                 stage1Data[2],
                                                 stage2Data[0],
                                                 stage2Data[1],
                                         }},
@@ skipping 146 matching lines @@
 
                                         Convey(fmt.Sprintf("stage %d", i), func() {
                                                 for j, expect := range stage.expect {
                                                         runner := func(f func(ic context.Context) error, _ *ds.TransactionOptions) error {
                                                                 return f(c)
                                                         }
                                                         if expect.inTxn {
                                                                 runner = data.RunInTransaction
                                                         }
 
+                                                        if expect.count == 0 {
+                                                                if len(expect.keys) > 0 {
+                                                                        expect.count = len(expect.keys)
+                                                                } else {
+                                                                        expect.count = len(expect.get)
+                                                                }
+                                                        }
+
                                                         if expect.keys != nil {
-                                                                err := runner(func(c context.Context) error {
-                                                                        data := ds.Get(c)
-                                                                        Convey(fmt.Sprintf("expect %d (keys)", j), func() {
+                                                                Convey(fmt.Sprintf("expect %d (keys)", j), func() {
+                                                                        err := runner(func(c context.Context) error {
+                                                                                data := ds.Get(c)
+                                                                                count, err := data.Count(expect.q)
+                                                                                So(err, ShouldBeNil)
+                                                                                So(count, ShouldEqual, expect.count)
+
                                                                                 rslt := []*ds.Key(nil)
                                                                                 So(data.GetAll(expect.q, &rslt), ShouldBeNil)
                                                                                 So(len(rslt), ShouldEqual, len(expect.keys))
                                                                                 for i, r := range rslt {
                                                                                         So(r, ShouldResemble, expect.keys[i])
                                                                                 }
-                                                                        })
-                                                                        return nil
-                                                                }, &ds.TransactionOptions{XG: true})
-                                                                So(err, ShouldBeNil)
+                                                                                return nil
+                                                                        }, &ds.TransactionOptions{XG: true})
+                                                                        So(err, ShouldBeNil)
+                                                                })
                                                         }
 
                                                         if expect.get != nil {
                                                                 Convey(fmt.Sprintf("expect %d (data)", j), func() {
                                                                         err := runner(func(c context.Context) error {
+                                                                                data := ds.Get(c)
+                                                                                count, err := data.Count(expect.q)
+                                                                                So(err, ShouldBeNil)
+                                                                                So(count, ShouldEqual, expect.count)
+
                                                                                 rslt := []ds.PropertyMap(nil)
                                                                                 So(data.GetAll(expect.q, &rslt), ShouldBeNil)
                                                                                 So(len(rslt), ShouldEqual, len(expect.get))
                                                                                 for i, r := range rslt {
                                                                                         So(r, ShouldResemble, expect.get[i])
                                                                                 }
                                                                                 return nil
                                                                         }, &ds.TransactionOptions{XG: true})
                                                                         So(err, ShouldBeNil)
                                                                 })
                                                         }
                                                 }
 
                                                 for j, fn := range stage.extraFns {
                                                         Convey(fmt.Sprintf("extraFn %d", j), func() {
                                                                 fn(c)
                                                         })
                                                 }
                                         })
                                 }
                         })
                 }
         })
 }