feat: replace the cosmos-db usecases in the tests with core/testing (#21525)

core/testing/memdb.go

@@ -3,6 +3,8 @@ package coretesting
 import (
 	"bytes"
 	"errors"
+	"fmt"
+	"sync"
 
 	"github.com/tidwall/btree"
 
@@ -15,7 +17,11 @@ const (
 	bTreeDegree = 32
 )
 
-var errKeyEmpty = errors.New("key cannot be empty")
+var (
+	errKeyEmpty    = errors.New("key cannot be empty")
+	errValueNil    = errors.New("value cannot be nil")
+	errBatchClosed = errors.New("batch is closed")
+)
 
 var _ store.KVStore = (*MemKV)(nil)
 
@@ -195,9 +201,6 @@ var errInvalidIterator = errors.New("invalid iterator")
 // If the iterator is not valid, it returns the errInvalidIterator error.
 // Otherwise, it returns nil.
 func (mi *memIterator) Error() error {
-	if !mi.Valid() {
-		return errInvalidIterator
-	}
 	return nil
 }
 
@@ -237,18 +240,187 @@ func (mi *memIterator) keyInRange(key []byte) bool {
 
 // Key returns the key of the current item in the iterator.
 func (mi *memIterator) Key() []byte {
+	mi.assertValid()
 	return mi.iter.Item().key
 }
 
 // Value returns the value of the current item in the iterator.
 func (mi *memIterator) Value() []byte {
+	mi.assertValid()
 	return mi.iter.Item().value
 }
 
 // assertValid checks if the memIterator is in a valid state.
 // If there is an error, it panics with the error message.
 func (mi *memIterator) assertValid() {
-	if err := mi.Error(); err != nil {
-		panic(err)
+	if !mi.valid {
+		panic(errInvalidIterator)
 	}
 }
+
+var _ store.KVStoreWithBatch = (*MemDB)(nil)
+
+// MemDB is a simple in-memory key-value store with Batch support.
+type MemDB struct {
+	kv  MemKV
+	mtx sync.RWMutex
+}
+
+// NewMemDB creates a new MemDB.
+func NewMemDB() store.KVStoreWithBatch {
+	return &MemDB{kv: NewMemKV()}
+}
+
+// KVStore implementation
+
+func (bt *MemDB) Get(key []byte) ([]byte, error) {
+	bt.mtx.RLock()
+	defer bt.mtx.RUnlock()
+	return bt.kv.Get(key)
+}
+
+func (bt *MemDB) Has(key []byte) (bool, error) {
+	bt.mtx.RLock()
+	defer bt.mtx.RUnlock()
+	return bt.kv.Has(key)
+}
+
+func (bt *MemDB) Set(key, value []byte) error {
+	bt.mtx.Lock()
+	defer bt.mtx.Unlock()
+	return bt.kv.Set(key, value)
+}
+
+func (bt *MemDB) Delete(key []byte) error {
+	bt.mtx.Lock()
+	defer bt.mtx.Unlock()
+	return bt.kv.Delete(key)
+}
+
+func (bt *MemDB) Iterator(start, end []byte) (store.Iterator, error) {
+	return bt.kv.Iterator(start, end)
+}
+
+func (bt *MemDB) ReverseIterator(start, end []byte) (store.Iterator, error) {
+	return bt.kv.ReverseIterator(start, end)
+}
+
+// Close closes the MemDB, releasing any resources held.
+func (db *MemDB) Close() error {
+	return nil
+}
+
+// NewBatch returns a new memDBBatch.
+func (db *MemDB) NewBatch() store.Batch {
+	return newMemDBBatch(db)
+}
+
+// NewBatchWithSize returns a new memDBBatch with the given size.
+func (db *MemDB) NewBatchWithSize(size int) store.Batch {
+	return newMemDBBatch(db)
+}
+
+// memDBBatch operations
+type opType int
+
+const (
+	opTypeSet opType = iota + 1
+	opTypeDelete
+)
+
+type operation struct {
+	opType
+	key   []byte
+	value []byte
+}
+
+// memDBBatch handles in-memory batching.
+type memDBBatch struct {
+	db   *MemDB
+	ops  []operation
+	size int
+}
+
+var _ store.Batch = (*memDBBatch)(nil)
+
+// newMemDBBatch creates a new memDBBatch
+func newMemDBBatch(db *MemDB) *memDBBatch {
+	return &memDBBatch{
+		db:   db,
+		ops:  []operation{},
+		size: 0,
+	}
+}
+
+// Set implements Batch.
+func (b *memDBBatch) Set(key, value []byte) error {
+	if len(key) == 0 {
+		return errKeyEmpty
+	}
+	if value == nil {
+		return errValueNil
+	}
+	if b.ops == nil {
+		return errBatchClosed
+	}
+	b.size += len(key) + len(value)
+	b.ops = append(b.ops, operation{opTypeSet, key, value})
+	return nil
+}
+
+// Delete implements Batch.
+func (b *memDBBatch) Delete(key []byte) error {
+	if len(key) == 0 {
+		return errKeyEmpty
+	}
+	if b.ops == nil {
+		return errBatchClosed
+	}
+	b.size += len(key)
+	b.ops = append(b.ops, operation{opTypeDelete, key, nil})
+	return nil
+}
+
+// Write implements Batch.
+func (b *memDBBatch) Write() error {
+	if b.ops == nil {
+		return errBatchClosed
+	}
+
+	b.db.mtx.Lock()
+	defer b.db.mtx.Unlock()
+
+	for _, op := range b.ops {
+		switch op.opType {
+		case opTypeSet:
+			b.db.kv.set(op.key, op.value)
+		case opTypeDelete:
+			b.db.kv.delete(op.key)
+		default:
+			return fmt.Errorf("unknown operation type %v (%v)", op.opType, op)
+		}
+	}
+
+	// Make sure batch cannot be used afterwards. Callers should still call Close(), for
+	return b.Close()
+}
+
+// WriteSync implements Batch.
+func (b *memDBBatch) WriteSync() error {
+	return b.Write()
+}
+
+// Close implements Batch.
+func (b *memDBBatch) Close() error {
+	b.ops = nil
+	b.size = 0
+	return nil
+}
+
+// GetByteSize implements Batch
+func (b *memDBBatch) GetByteSize() (int, error) {
+	if b.ops == nil {
+		return 0, errBatchClosed
+	}
+	return b.size, nil
+}