From 7909c1837ea6a54ed13d69f2f70e5c4ed4a173aa Mon Sep 17 00:00:00 2001
From: yihuang <huang@crypto.com>
Date: Fri, 8 Sep 2023 13:20:20 +0800
Subject: [PATCH] Problem: no unit test for native action (#337)

* Problem: no unit test for native action

Solution:
- add unit test

* test nested cases
---
 x/evm/statedb/mock_test.go    | 16 +++++--
 x/evm/statedb/statedb_test.go | 88 +++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 3 deletions(-)

diff --git a/x/evm/statedb/mock_test.go b/x/evm/statedb/mock_test.go
index abcef5b4ce..311fed6b88 100644
--- a/x/evm/statedb/mock_test.go
+++ b/x/evm/statedb/mock_test.go
@@ -26,17 +26,23 @@ type MockAcount struct {
 type MockKeeper struct {
 	accounts map[common.Address]MockAcount
 	codes    map[common.Hash][]byte
+	keys     map[string]*storetypes.KVStoreKey
 }
 
-func NewMockKeeper() *MockKeeper {
+func NewMockKeeperWithKeys(keys map[string]*storetypes.KVStoreKey) *MockKeeper {
 	return &MockKeeper{
 		accounts: make(map[common.Address]MockAcount),
 		codes:    make(map[common.Hash][]byte),
+		keys:     keys,
 	}
 }
 
+func NewMockKeeper() *MockKeeper {
+	return NewMockKeeperWithKeys(nil)
+}
+
 func (k MockKeeper) StoreKeys() map[string]*storetypes.KVStoreKey {
-	return nil
+	return k.keys
 }
 
 func (k MockKeeper) GetAccount(ctx sdk.Context, addr common.Address) *statedb.Account {
@@ -115,5 +121,9 @@ func (k MockKeeper) Clone() *MockKeeper {
 	for k, v := range k.codes {
 		codes[k] = v
 	}
-	return &MockKeeper{accounts, codes}
+	keys := make(map[string]*storetypes.KVStoreKey, len(k.keys))
+	for k, v := range k.keys {
+		keys[k] = v
+	}
+	return &MockKeeper{accounts, codes, keys}
 }
diff --git a/x/evm/statedb/statedb_test.go b/x/evm/statedb/statedb_test.go
index 3a491aa8cc..a427dbe75b 100644
--- a/x/evm/statedb/statedb_test.go
+++ b/x/evm/statedb/statedb_test.go
@@ -1,9 +1,15 @@
 package statedb_test
 
 import (
+	"errors"
 	"math/big"
 	"testing"
 
+	dbm "github.com/cometbft/cometbft-db"
+	"github.com/cometbft/cometbft/libs/log"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	"github.com/cosmos/cosmos-sdk/store/rootmulti"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/ethereum/go-ethereum/common"
 	ethtypes "github.com/ethereum/go-ethereum/core/types"
@@ -565,6 +571,88 @@ func (suite *StateDBTestSuite) TestIterateStorage() {
 	suite.Require().Equal(1, len(storage))
 }
 
+func (suite *StateDBTestSuite) TestNativeAction() {
+	db := dbm.NewMemDB()
+	ms := rootmulti.NewStore(db, log.NewNopLogger())
+	keys := map[string]*storetypes.KVStoreKey{
+		"storekey": storetypes.NewKVStoreKey("storekey"),
+	}
+	ms.MountStoreWithDB(keys["storekey"], storetypes.StoreTypeIAVL, nil)
+	suite.Require().NoError(ms.LoadLatestVersion())
+	ctx := sdk.NewContext(ms, tmproto.Header{}, false, log.NewNopLogger())
+
+	keeper := NewMockKeeperWithKeys(keys)
+	stateDB := statedb.New(ctx, keeper, emptyTxConfig)
+
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		store.Set([]byte("success1"), []byte("value"))
+		return nil
+	})
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		store.Set([]byte("failure1"), []byte("value"))
+		return errors.New("failure")
+	})
+
+	// test query
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		suite.Require().Equal([]byte("value"), store.Get([]byte("success1")))
+		suite.Require().Nil(store.Get([]byte("failure1")))
+		return nil
+	})
+
+	rev1 := stateDB.Snapshot()
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		store.Set([]byte("success2"), []byte("value"))
+		return nil
+	})
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		store.Set([]byte("failure2"), []byte("value"))
+		return errors.New("failure")
+	})
+
+	// test query
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		suite.Require().Equal([]byte("value"), store.Get([]byte("success1")))
+		suite.Require().Equal([]byte("value"), store.Get([]byte("success2")))
+		suite.Require().Nil(store.Get([]byte("failure2")))
+		return nil
+	})
+
+	stateDB.RevertToSnapshot(rev1)
+
+	_ = stateDB.Snapshot()
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		store.Set([]byte("success3"), []byte("value"))
+		return nil
+	})
+
+	// test query
+	stateDB.ExecuteNativeAction(func(ctx sdk.Context) error {
+		store := ctx.KVStore(keys["storekey"])
+		suite.Require().Equal([]byte("value"), store.Get([]byte("success1")))
+		suite.Require().Nil(store.Get([]byte("success2")))
+		suite.Require().Equal([]byte("value"), store.Get([]byte("success3")))
+		return nil
+	})
+
+	suite.Require().NoError(stateDB.Commit())
+
+	// query committed state
+	store := ctx.KVStore(keys["storekey"])
+	suite.Require().Equal([]byte("value"), store.Get([]byte("success1")))
+	suite.Require().Nil(store.Get([]byte("success2")))
+	suite.Require().Equal([]byte("value"), store.Get([]byte("success3")))
+	suite.Require().Nil(store.Get([]byte("failure1")))
+	suite.Require().Nil(store.Get([]byte("failure2")))
+}
+
 func CollectContractStorage(db vm.StateDB) statedb.Storage {
 	storage := make(statedb.Storage)
 	db.ForEachStorage(address, func(k, v common.Hash) bool {