testutils: Add script-based test utilities

Add generic testscript commands for testing against
StateDB tables.

This allows implementing tests as scripts, which becomes useful when
tests perform multiple steps on tables and need to verify the output
each step.

Signed-off-by: Jussi Maki <[email protected]>

any_table.go

@@ -0,0 +1,63 @@
+package statedb
+
+import (
+	"iter"
+)
+
+// AnyTable allows any-typed access to a StateDB table. This is intended
+// for building generic tooling for accessing the table and should be
+// avoided if possible.
+type AnyTable struct {
+	Meta TableMeta
+}
+
+func (t AnyTable) All(txn ReadTxn) iter.Seq2[any, Revision] {
+	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
+	return anySeq(indexTxn.Iterator())
+}
+
+func (t AnyTable) UnmarshalYAML(data []byte) (any, error) {
+	return t.Meta.unmarshalYAML(data)
+}
+
+func (t AnyTable) Insert(txn WriteTxn, obj any) (old any, hadOld bool, err error) {
+	var iobj object
+	iobj, hadOld, err = txn.getTxn().insert(t.Meta, Revision(0), obj)
+	if hadOld {
+		old = iobj.data
+	}
+	return
+}
+
+func (t AnyTable) Delete(txn WriteTxn, obj any) (old any, hadOld bool, err error) {
+	var iobj object
+	iobj, hadOld, err = txn.getTxn().delete(t.Meta, Revision(0), obj)
+	if hadOld {
+		old = iobj.data
+	}
+	return
+}
+
+func (t AnyTable) Prefix(txn ReadTxn, key string) iter.Seq2[any, Revision] {
+	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
+	iter, _ := indexTxn.Prefix([]byte(key))
+	return anySeq(iter)
+}
+
+func (t AnyTable) LowerBound(txn ReadTxn, key string) iter.Seq2[any, Revision] {
+	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
+	iter := indexTxn.LowerBound([]byte(key))
+	return anySeq(iter)
+}
+
+func (t AnyTable) TableHeader() []string {
+	zero := t.Meta.proto()
+	if tw, ok := zero.(TableWritable); ok {
+		return tw.TableHeader()
+	}
+	return nil
+}
+
+func (t AnyTable) Proto() any {
+	return t.Meta.proto()
+}

db.go

@@ -247,6 +247,25 @@ func (db *DB) WriteTxn(table TableMeta, tables ...TableMeta) WriteTxn {
 	return txn
 }
 
+func (db *DB) GetTables(txn ReadTxn) (tbls []TableMeta) {
+	root := txn.getTxn().root
+	tbls = make([]TableMeta, 0, len(root))
+	for _, table := range root {
+		tbls = append(tbls, table.meta)
+	}
+	return
+}
+
+func (db *DB) GetTable(txn ReadTxn, name string) TableMeta {
+	root := txn.getTxn().root
+	for _, table := range root {
+		if table.meta.Name() == name {
+			return table.meta
+		}
+	}
+	return nil
+}
+
 // Start the background workers for the database.
 //
 // This starts the graveyard worker that deals with garbage collecting

go.mod

@@ -5,11 +5,13 @@ go 1.23
 require (
 	github.com/cilium/hive v0.0.0-20240209163124-bd6ebb4ec11d
 	github.com/cilium/stream v0.0.0-20240209152734-a0792b51812d
+	github.com/rogpeppe/go-internal v1.11.0
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.8.4
 	go.uber.org/goleak v1.3.0
 	golang.org/x/time v0.5.0
+	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
@@ -34,7 +36,7 @@ require (
 	golang.org/x/sys v0.17.0 // indirect
 	golang.org/x/term v0.16.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/tools v0.17.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -77,6 +77,8 @@ golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
 golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
+golang.org/x/tools v0.17.0 h1:FvmRgNOcs3kOa+T20R1uhfP9F6HgG2mfxDv1vrx1Htc=
+golang.org/x/tools v0.17.0/go.mod h1:xsh6VxdV005rRVaS6SSAf9oiAqljS7UZUacMZ8Bnsps=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

iterator.go

@@ -53,7 +53,7 @@ func ToSeq[A, B any](seq iter.Seq2[A, B]) iter.Seq[A] {
 	}
 }
 
-// partSeq returns a sequence of objects from a part Iterator.
+// partSeq returns a casted sequence of objects from a part Iterator.
 func partSeq[Obj any](iter *part.Iterator[object]) iter.Seq2[Obj, Revision] {
 	return func(yield func(Obj, Revision) bool) {
 		// Iterate over a clone of the original iterator to allow the sequence to be iterated
@@ -71,6 +71,24 @@ func partSeq[Obj any](iter *part.Iterator[object]) iter.Seq2[Obj, Revision] {
 	}
 }
 
+// anySeq returns a sequence of objects from a part Iterator.
+func anySeq(iter *part.Iterator[object]) iter.Seq2[any, Revision] {
+	return func(yield func(any, Revision) bool) {
+		// Iterate over a clone of the original iterator to allow the sequence to be iterated
+		// from scratch multiple times.
+		it := iter.Clone()
+		for {
+			_, iobj, ok := it.Next()
+			if !ok {
+				break
+			}
+			if !yield(iobj.data, iobj.revision) {
+				break
+			}
+		}
+	}
+}
+
 // nonUniqueSeq returns a sequence of objects for a non-unique index.
 // Non-unique indexes work by concatenating the secondary key with the
 // primary key and then prefix searching for the items:
@@ -128,6 +146,13 @@ func (it *iterator[Obj]) Next() (obj Obj, revision uint64, ok bool) {
 	return
 }
 
+// Iterator for iterating a sequence objects.
+type Iterator[Obj any] interface {
+	// Next returns the next object and its revision if ok is true, otherwise
+	// zero values to mean that the iteration has finished.
+	Next() (obj Obj, rev Revision, ok bool)
+}
+
 func NewDualIterator[Obj any](left, right Iterator[Obj]) *DualIterator[Obj] {
 	return &DualIterator[Obj]{
 		left:  iterState[Obj]{iter: left},

table.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/cilium/statedb/internal"
 	"github.com/cilium/statedb/part"
+	"gopkg.in/yaml.v3"
 
 	"github.com/cilium/statedb/index"
 )
@@ -184,6 +185,15 @@ func (t *genTable[Obj]) Name() string {
 	return t.table
 }
 
+func (t *genTable[Obj]) Indexes() []string {
+	idxs := make([]string, 0, 1+len(t.secondaryAnyIndexers))
+	idxs = append(idxs, t.primaryAnyIndexer.name)
+	for k := range t.secondaryAnyIndexers {
+		idxs = append(idxs, k)
+	}
+	return idxs
+}
+
 func (t *genTable[Obj]) ToTable() Table[Obj] {
 	return t
 }
@@ -468,5 +478,18 @@ func (t *genTable[Obj]) sortableMutex() internal.SortableMutex {
 	return t.smu
 }
 
+func (t *genTable[Obj]) proto() any {
+	var zero Obj
+	return zero
+}
+
+func (t *genTable[Obj]) unmarshalYAML(data []byte) (any, error) {
+	var obj Obj
+	if err := yaml.Unmarshal(data, &obj); err != nil {
+		return nil, err
+	}
+	return obj, nil
+}
+
 var _ Table[bool] = &genTable[bool]{}
 var _ RWTable[bool] = &genTable[bool]{}