[WIP] persistence context rollback tests

persistence/context.go

@@ -81,6 +81,11 @@ func (p *PostgresContext) Commit(proposerAddr, quorumCert []byte) error {
 		return err
 	}
 
+	// Commit the tree store updates. This is the grossness I'm talking about.
+	if err := p.stateTrees.Commit(); err != nil {
+		return err
+	}
+
 	// Commit the SQL transaction
 	ctx := context.TODO()
 	if err := p.tx.Commit(ctx); err != nil {

persistence/test/state_test.go

@@ -187,8 +187,103 @@ func TestStateHash_ReplayingRandomTransactionsIsDeterministic(t *testing.T) {
 	}
 }
 
+func TestStateHash_TreeUpdatesRollback(t *testing.T) {
+	height := int64(0)
+	db := NewTestPostgresContext(t, height)
+
+	apps, err := db.GetAllApps(height)
+	require.NoError(t, err)
+	app := apps[0]
+
+	addrBz, err := hex.DecodeString(app.GetAddress())
+	require.NoError(t, err)
+	require.NotEmpty(t, addrBz)
+
+	hash1, err := db.ComputeStateHash()
+	require.NoError(t, err)
+
+	err = db.NewSavePoint()
+	require.NoError(t, err)
+
+	// Update the app's stake to trigger a change when ComputeStateHash is called
+	stakeAmount := 42
+	stakeAmount += 1 // change the stake amount
+	stakeAmountStr := strconv.Itoa(stakeAmount)
+	err = db.SetAppStakeAmount(addrBz, stakeAmountStr)
+	require.NoError(t, err)
+
+	txBz := []byte("a tx, i am, which set the app stake amount to " + stakeAmountStr)
+	idxTx := &coreTypes.IndexedTransaction{
+		Tx:            txBz,
+		Height:        height,
+		Index:         0,
+		ResultCode:    0,
+		Error:         "TODO",
+		SignerAddr:    "TODO",
+		RecipientAddr: "TODO",
+		MessageType:   "TODO",
+	}
+
+	err = db.IndexTransaction(idxTx)
+	require.NoError(t, err)
+
+	hash2, err := db.ComputeStateHash()
+	require.NoError(t, err)
+	require.NotEmpty(t, hash2)
+
+	err = db.RollbackToSavePoint()
+	require.NoError(t, err)
+
+	hash3, err := db.ComputeStateHash()
+	require.NoError(t, err)
+	require.Equal(t, hash3, hash1)
+}
+
 func TestStateHash_TreeUpdatesAreIdempotent(t *testing.T) {
-	// ADDTEST(#361): Create an issue dedicated to increasing the test coverage for state hashes
+	height := int64(0)
+	db := NewTestPostgresContext(t, height)
+
+	apps, err := db.GetAllApps(height)
+	require.NoError(t, err)
+	app := apps[0]
+
+	addrBz, err := hex.DecodeString(app.GetAddress())
+	require.NoError(t, err)
+	require.NotEmpty(t, addrBz)
+
+	err = db.NewSavePoint()
+	require.NoError(t, err)
+
+	// Update the app's stake
+	stakeAmount := 42
+	stakeAmount += 1 // change the stake amount
+	stakeAmountStr := strconv.Itoa(stakeAmount)
+	err = db.SetAppStakeAmount(addrBz, stakeAmountStr)
+	require.NoError(t, err)
+
+	txBz := []byte("a tx, i am, which set the app stake amount to " + stakeAmountStr)
+	idxTx := &coreTypes.IndexedTransaction{
+		Tx:            txBz,
+		Height:        height,
+		Index:         0,
+		ResultCode:    0,
+		Error:         "TODO",
+		SignerAddr:    "TODO",
+		RecipientAddr: "TODO",
+		MessageType:   "TODO",
+	}
+
+	err = db.IndexTransaction(idxTx)
+	require.NoError(t, err)
+
+	hash1, err := db.ComputeStateHash()
+	require.NoError(t, err)
+	require.NotEmpty(t, hash1)
+
+	hash2, err := db.ComputeStateHash()
+	require.NoError(t, err)
+	require.NotEmpty(t, hash1)
+	require.Equal(t, hash1, hash2)
 }
 
 func TestStateHash_TreeUpdatesNegativeTestCase(t *testing.T) {

persistence/trees/atomic_test.go

@@ -0,0 +1,21 @@
+package trees
+
+import (
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	mockModules "github.com/pokt-network/pocket/shared/modules/mocks"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTreeStore_AtomicUpdates(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockBus := mockModules.NewMockBus(ctrl)
+	ts := &treeStore{
+		bus:          mockBus,
+		txi:          mockBus.GetPersistenceModule().GetTxIndexer(),
+		treeStoreDir: ":memory:",
+	}
+	ts.setupTrees()
+	require.NotEmpty(t, ts)
+}

persistence/trees/module.go

@@ -9,19 +9,19 @@ import (
 	"github.com/pokt-network/smt"
 )
 
-var _ modules.Module = &treeStore{}
+var _ modules.Module = &TreeStore{}
 
-func (*treeStore) Create(bus modules.Bus, options ...modules.ModuleOption) (modules.Module, error) {
-	m := &treeStore{}
+func (*TreeStore) Create(bus modules.Bus, options ...modules.ModuleOption) (modules.Module, error) {
+	m := &TreeStore{}
 
 	bus.RegisterModule(m)
 
 	for _, option := range options {
 		option(m)
 	}
 
-	if m.txi == nil {
-		m.txi = bus.GetPersistenceModule().GetTxIndexer()
+	if m.TXI == nil {
+		m.TXI = bus.GetPersistenceModule().GetTxIndexer()
 	}
 
 	if err := m.setupTrees(); err != nil {
@@ -32,13 +32,13 @@ func (*treeStore) Create(bus modules.Bus, options ...modules.ModuleOption) (modu
 }
 
 func Create(bus modules.Bus, options ...modules.ModuleOption) (modules.Module, error) {
-	return new(treeStore).Create(bus, options...)
+	return new(TreeStore).Create(bus, options...)
 }
 
 // WithLogger assigns a logger for the tree store
 func WithLogger(logger *modules.Logger) modules.ModuleOption {
 	return func(m modules.InitializableModule) {
-		if mod, ok := m.(*treeStore); ok {
+		if mod, ok := m.(*TreeStore); ok {
 			mod.logger = logger
 		}
 	}
@@ -48,35 +48,35 @@ func WithLogger(logger *modules.Logger) modules.ModuleOption {
 // saves its data.
 func WithTreeStoreDirectory(path string) modules.ModuleOption {
 	return func(m modules.InitializableModule) {
-		mod, ok := m.(*treeStore)
+		mod, ok := m.(*TreeStore)
 		if ok {
-			mod.treeStoreDir = path
+			mod.TreeStoreDir = path
 		}
 	}
 }
 
 // WithTxIndexer assigns a TxIndexer for use during operation.
 func WithTxIndexer(txi indexer.TxIndexer) modules.ModuleOption {
 	return func(m modules.InitializableModule) {
-		mod, ok := m.(*treeStore)
+		mod, ok := m.(*TreeStore)
 		if ok {
-			mod.txi = txi
+			mod.TXI = txi
 		}
 	}
 }
 
-func (t *treeStore) GetModuleName() string  { return modules.TreeStoreModuleName }
-func (t *treeStore) Start() error           { return nil }
-func (t *treeStore) Stop() error            { return nil }
-func (t *treeStore) GetBus() modules.Bus    { return t.bus }
-func (t *treeStore) SetBus(bus modules.Bus) { t.bus = bus }
+func (t *TreeStore) GetModuleName() string  { return modules.TreeStoreModuleName }
+func (t *TreeStore) Start() error           { return nil }
+func (t *TreeStore) Stop() error            { return nil }
+func (t *TreeStore) GetBus() modules.Bus    { return t.Bus }
+func (t *TreeStore) SetBus(bus modules.Bus) { t.Bus = bus }
 
-func (t *treeStore) setupTrees() error {
-	if t.treeStoreDir == ":memory:" {
+func (t *TreeStore) setupTrees() error {
+	if t.TreeStoreDir == ":memory:" {
 		return t.setupInMemory()
 	}
 
-	nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.treeStoreDir, RootTreeName))
+	nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.TreeStoreDir, RootTreeName))
 	if err != nil {
 		return err
 	}
@@ -88,7 +88,7 @@ func (t *treeStore) setupTrees() error {
 	t.merkleTrees = make(map[string]*stateTree, len(stateTreeNames))
 
 	for i := 0; i < len(stateTreeNames); i++ {
-		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.treeStoreDir, stateTreeNames[i]))
+		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.TreeStoreDir, stateTreeNames[i]))
 		if err != nil {
 			return err
 		}
@@ -102,7 +102,7 @@ func (t *treeStore) setupTrees() error {
 	return nil
 }
 
-func (t *treeStore) setupInMemory() error {
+func (t *TreeStore) setupInMemory() error {
 	nodeStore := kvstore.NewMemKVStore()
 	t.rootTree = &stateTree{
 		name:      RootTreeName,