treapstore: Add Collection.VisitItemsAscend.

common/data/treapstore/store_test.go

 // Copyright 2016 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.
 
 package treapstore
 
 import (
         "fmt"
         "strconv"
         "strings"
         "testing"
 
         "github.com/luci/gtreap"
 
         . "github.com/smartystreets/goconvey/convey"
 )
 
 func stringCompare(a, b interface{}) int {
         return strings.Compare(a.(string), b.(string))
 }
 
 func putMulti(c *Collection, vs ...string) {
         for _, v := range vs {
                 c.Put(v)
         }
 }
 
+func visitAll(c *Collection, pivot string) []string {
+        res := []string{}
+        c.VisitAscend(pivot, func(v gtreap.Item) bool {
+                res = append(res, v.(string))
+                return true
+        })
+        return res
+}
+
 func iterAll(it *gtreap.Iterator) []string {
         all := []string{}
         for {
                 v, ok := it.Next()
                 if !ok {
                         return all
                 }
                 all = append(all, v.(string))
         }
 }
 
+func shouldHaveKeys(actual interface{}, expected ...interface{}) string {
+        c := actual.(*Collection)
+
+        // expected can either be a single []string or a series of strings.
+        var keys []string
+        var ok bool
+        if len(expected) == 1 {
+                keys, ok = expected[0].([]string)
+        }
+        if !ok {
+                keys = make([]string, len(expected))
+                for i, v := range expected {
+                        keys[i] = v.(string)
+                }
+        }
+
+        if err := ShouldResemble(iterAll(c.Iterator("")), keys); err != "" {
+                return fmt.Sprintf("failed via iterator: %s", err)
+        }
+        if err := ShouldResemble(visitAll(c, ""), keys); err != "" {
+                return fmt.Sprintf("failed via visit: %s", err)
+        }
+        return ""
+}
+
 func TestStore(t *testing.T) {
         t.Parallel()
 
         Convey(`Testing a string Store`, t, func() {
                 st := New()
                 coll := st.CreateCollection("test", stringCompare)
 
                 Convey(`When empty`, func() {
                         checkEmpty := func(c *Collection) {
                                 So(c.Get("foo"), ShouldBeNil)
-                                So(iterAll(c.Iterator("")), ShouldResemble, []string{})
+                                So(c, shouldHaveKeys)
                         }
 
                         // Check the basic Store.
                         checkEmpty(coll)
 
                         // Take a snapshot, then mutate the base Store. Assert that the snapshot
                         // is still empty.
                         snap := st.Snapshot()
                         coll.Put("foo")
                         checkEmpty(snap.GetCollection("test"))
                 })
 
                 Convey(`With keys`, func() {
                         putMulti(coll, "x", "w", "b", "a")
 
                         Convey(`Can iterate`, func() {
                                 checkKeys := func(coll *Collection, keys ...string) {
                                         for _, k := range keys {
                                                 So(coll.Get(k), ShouldEqual, k)
                                         }
 
-                                        So(iterAll(coll.Iterator("")), ShouldResemble, keys)
+                                        So(coll, shouldHaveKeys, keys)
                                         for i, k := range keys {
                                                 So(iterAll(coll.Iterator(k)), ShouldResemble, keys[i:])
                                                 So(iterAll(coll.Iterator(k+"1")), ShouldResemble, keys[i+1:])
                                         }
                                 }
                                 checkKeys(coll, "a", "b", "w", "x")
 
                                 // Take a snapshot, then mutate the base Store. Assert that the snapshot
                                 // is still empty.
                                 snap := st.Snapshot()
@@ skipping 184 matching lines @@
                                 // The first time around, we *could* read before anything has been
                                 // written. Every other time, something from the previous round will
                                 // have been written.
                                 check = ShouldBeGreaterThanOrEqualTo
                         }
                         So(<-readDoneC, check, 0)
                         snaps = append(snaps, head.Snapshot())
                 }
         })
 }