chore: move common function to keeper_test.go

piplabs · Sep 5, 2024 · 2b4e09a · 2b4e09a
1 parent d45cb9e
commit 2b4e09a
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 76 deletions.
diff --git a/client/x/evmstaking/keeper/keeper_test.go b/client/x/evmstaking/keeper/keeper_test.go
@@ -1,6 +1,7 @@
 package keeper_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -9,6 +10,7 @@ import (
 	"cosmossdk.io/store/metrics"
 	storetypes "cosmossdk.io/store/types"
 
+	"github.com/cometbft/cometbft/crypto"
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	dbm "github.com/cosmos/cosmos-db"
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -21,6 +23,7 @@ import (
 	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	skeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	"github.com/cosmos/cosmos-sdk/x/staking/testutil"
 	stypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/stretchr/testify/suite"
 
@@ -29,6 +32,7 @@ import (
 	estestutil "github.com/piplabs/story/client/x/evmstaking/testutil"
 	"github.com/piplabs/story/client/x/evmstaking/types"
 	"github.com/piplabs/story/lib/ethclient"
+	"github.com/piplabs/story/lib/k1util"
 
 	"go.uber.org/mock/gomock"
 )
@@ -134,6 +138,39 @@ func TestTestSuite(t *testing.T) {
 	suite.Run(t, new(TestSuite))
 }
 
+// setupValidatorAndDelegation creates a validator and delegation for testing.
+func (s *TestSuite) setupValidatorAndDelegation(ctx context.Context, valPubKey, delPubKey crypto.PubKey, valAddr sdk.ValAddress, delAddr sdk.AccAddress) {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	require := s.Require()
+	bankKeeper, stakingKeeper, keeper := s.BankKeeper, s.StakingKeeper, s.EVMStakingKeeper
+
+	// Convert public key to cosmos format
+	valCosmosPubKey, err := k1util.PubKeyToCosmos(valPubKey)
+	require.NoError(err)
+
+	// Create and update validator
+	val := testutil.NewValidator(s.T(), valAddr, valCosmosPubKey)
+	valTokens := stakingKeeper.TokensFromConsensusPower(ctx, 10)
+	validator, _ := val.AddTokensFromDel(valTokens)
+	bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), stypes.NotBondedPoolName, stypes.BondedPoolName, gomock.Any())
+	_ = skeeper.TestingUpdateValidator(stakingKeeper, sdkCtx, validator, true)
+
+	// Create and set delegation
+	delAmt := stakingKeeper.TokensFromConsensusPower(ctx, 100).ToLegacyDec()
+	delegation := stypes.NewDelegation(delAddr.String(), valAddr.String(), delAmt)
+	require.NoError(stakingKeeper.SetDelegation(ctx, delegation))
+
+	// Map delegator to EVM address
+	delEvmAddr, err := k1util.CosmosPubkeyToEVMAddress(delPubKey.Bytes())
+	require.NoError(err)
+	require.NoError(keeper.DelegatorMap.Set(ctx, delAddr.String(), delEvmAddr.String()))
+
+	// Ensure delegation is set correctly
+	delegation, err = stakingKeeper.GetDelegation(ctx, delAddr, valAddr)
+	require.NoError(err)
+	require.Equal(delAmt, delegation.GetShares())
+}
+
 func createCorruptedPubKey(pubKey []byte) []byte {
 	corruptedPubKey := append([]byte(nil), pubKey...)
 	corruptedPubKey[0] = 0x04

diff --git a/client/x/evmstaking/keeper/msg_server_test.go b/client/x/evmstaking/keeper/msg_server_test.go
@@ -3,51 +3,11 @@ package keeper_test
 import (
 	"context"
 
-	"github.com/cometbft/cometbft/crypto"
-	sdk "github.com/cosmos/cosmos-sdk/types"
-	skeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
-	"github.com/cosmos/cosmos-sdk/x/staking/testutil"
 	stypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	"github.com/piplabs/story/client/x/evmstaking/types"
-	"github.com/piplabs/story/lib/k1util"
-
-	"go.uber.org/mock/gomock"
 )
 
-// setupValidatorAndDelegation creates a validator and delegation for testing.
-func (s *TestSuite) setupValidatorAndDelegation(ctx context.Context, valPubKey, delPubKey crypto.PubKey, valAddr sdk.ValAddress, delAddr sdk.AccAddress) {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	require := s.Require()
-	bankKeeper, stakingKeeper, keeper := s.BankKeeper, s.StakingKeeper, s.EVMStakingKeeper
-
-	// Convert public key to cosmos format
-	valCosmosPubKey, err := k1util.PubKeyToCosmos(valPubKey)
-	require.NoError(err)
-
-	// Create and update validator
-	val := testutil.NewValidator(s.T(), valAddr, valCosmosPubKey)
-	valTokens := stakingKeeper.TokensFromConsensusPower(ctx, 10)
-	validator, _ := val.AddTokensFromDel(valTokens)
-	bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), stypes.NotBondedPoolName, stypes.BondedPoolName, gomock.Any())
-	_ = skeeper.TestingUpdateValidator(stakingKeeper, sdkCtx, validator, true)
-
-	// Create and set delegation
-	delAmt := stakingKeeper.TokensFromConsensusPower(ctx, 100).ToLegacyDec()
-	delegation := stypes.NewDelegation(delAddr.String(), valAddr.String(), delAmt)
-	require.NoError(stakingKeeper.SetDelegation(ctx, delegation))
-
-	// Map delegator to EVM address
-	delEvmAddr, err := k1util.CosmosPubkeyToEVMAddress(delPubKey.Bytes())
-	require.NoError(err)
-	require.NoError(keeper.DelegatorMap.Set(ctx, delAddr.String(), delEvmAddr.String()))
-
-	// Ensure delegation is set correctly
-	delegation, err = stakingKeeper.GetDelegation(ctx, delAddr, valAddr)
-	require.NoError(err)
-	require.Equal(delAmt, delegation.GetShares())
-}
-
 func (s *TestSuite) TestAddWithdrawal() {
 	require := s.Require()
 	ctx, msgServer, keeper := s.Ctx, s.msgServer, s.EVMStakingKeeper

diff --git a/client/x/evmstaking/keeper/withdraw_test.go b/client/x/evmstaking/keeper/withdraw_test.go
@@ -7,11 +7,8 @@ import (
 
 	sdkmath "cosmossdk.io/math"
 
-	"github.com/cometbft/cometbft/crypto"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	dtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
-	skeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
-	"github.com/cosmos/cosmos-sdk/x/staking/testutil"
 	stypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/ethereum/go-ethereum/common"
 	gethtypes "github.com/ethereum/go-ethereum/core/types"
@@ -25,39 +22,6 @@ import (
 	"go.uber.org/mock/gomock"
 )
 
-// setupValidatorAndDelegation creates a validator and delegation for testing.
-func (s *TestSuite) setupValidatorAndDelegation(ctx sdk.Context, valPubKey, delPubKey crypto.PubKey, valAddr sdk.ValAddress, delAddr sdk.AccAddress) {
-	require := s.Require()
-	stakingKeeper := s.StakingKeeper
-	keeper := s.EVMStakingKeeper
-
-	// Convert public key to cosmos format
-	valCosmosPubKey, err := k1util.PubKeyToCosmos(valPubKey)
-	require.NoError(err)
-
-	// Create and update validator
-	val := testutil.NewValidator(s.T(), valAddr, valCosmosPubKey)
-	valTokens := stakingKeeper.TokensFromConsensusPower(ctx, 10)
-	validator, _ := val.AddTokensFromDel(valTokens)
-	s.BankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), stypes.NotBondedPoolName, stypes.BondedPoolName, gomock.Any())
-	_ = skeeper.TestingUpdateValidator(stakingKeeper, ctx, validator, true)
-
-	// Create and set delegation
-	delAmt := stakingKeeper.TokensFromConsensusPower(ctx, 100).ToLegacyDec()
-	delegation := stypes.NewDelegation(delAddr.String(), valAddr.String(), delAmt)
-	require.NoError(stakingKeeper.SetDelegation(ctx, delegation))
-
-	// Map delegator to EVM address
-	delEvmAddr, err := k1util.CosmosPubkeyToEVMAddress(delPubKey.Bytes())
-	require.NoError(err)
-	require.NoError(keeper.DelegatorMap.Set(ctx, delAddr.String(), delEvmAddr.String()))
-
-	// Ensure delegation is set correctly
-	delegation, err = stakingKeeper.GetDelegation(ctx, delAddr, valAddr)
-	require.NoError(err)
-	require.Equal(delAmt, delegation.GetShares())
-}
-
 func (s *TestSuite) TestExpectedPartialWithdrawals() {
 	require := s.Require()
 	ctx, keeper, stakingKeeper, distrKeeper := s.Ctx, s.EVMStakingKeeper, s.StakingKeeper, s.DistrKeeper