swarm/state: refactor InmemoryStore (#18143)

1 files changed, 21 insertions, 0 deletions

diff --git a/swarm/state/dbstore.go b/swarm/state/dbstore.go
index b0aa92e27..fc5dd8f7c 100644
--- a/swarm/state/dbstore.go
+++ b/swarm/state/dbstore.go
@@ -22,6 +22,7 @@ import (
 	"errors"
 
 	"github.com/syndtr/goleveldb/leveldb"
+	"github.com/syndtr/goleveldb/leveldb/storage"
 )
 
 // ErrNotFound is returned when no results are returned from the database
@@ -30,6 +31,15 @@ var ErrNotFound = errors.New("ErrorNotFound")
 // ErrInvalidArgument is returned when the argument type does not match the expected type
 var ErrInvalidArgument = errors.New("ErrorInvalidArgument")
 
+// Store defines methods required to get, set, delete values for different keys
+// and close the underlying resources.
+type Store interface {
+	Get(key string, i interface{}) (err error)
+	Put(key string, i interface{}) (err error)
+	Delete(key string) (err error)
+	Close() error
+}
+
 // DBStore uses LevelDB to store values.
 type DBStore struct {
 	db *leveldb.DB
@@ -46,6 +56,17 @@ func NewDBStore(path string) (s *DBStore, err error) {
 	}, nil
 }
 
+// NewInmemoryStore returns a new instance of DBStore. To be used only in tests and simulations.
+func NewInmemoryStore() *DBStore {
+	db, err := leveldb.Open(storage.NewMemStorage(), nil)
+	if err != nil {
+		panic(err)
+	}
+	return &DBStore{
+		db: db,
+	}
+}
+
 // Get retrieves a persisted value for a specific key. If there is no results
 // ErrNotFound is returned. The provided parameter should be either a byte slice or
 // a struct that implements the encoding.BinaryUnmarshaler interface