add checking authority and module account

Author: nkitlabs

x/bandtss/genesis.go

@@ -1,6 +1,8 @@
 package bandtss
 
 import (
+	"fmt"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	"github.com/bandprotocol/chain/v2/x/bandtss/keeper"
@@ -9,6 +11,17 @@ import (
 
 // InitGenesis performs genesis initialization for this module.
 func InitGenesis(ctx sdk.Context, k *keeper.Keeper, data *types.GenesisState) {
+	// check if the module account exists
+	moduleAcc := k.GetBandtssAccount(ctx)
+	if moduleAcc == nil {
+		panic(fmt.Sprintf("%s module account has not been set", types.ModuleName))
+	}
+
+	// Set module account if its balance is zero
+	if balance := k.GetModuleBalance(ctx); balance.IsZero() {
+		k.SetModuleAccount(ctx, moduleAcc)
+	}
+
 	if err := k.SetParams(ctx, data.Params); err != nil {
 		panic(err)
 	}

x/bandtss/genesis_test.go

@@ -3,7 +3,10 @@ package bandtss_test
 import (
 	"testing"
 
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 
 	"github.com/bandprotocol/chain/v2/pkg/tss"
 	bandtesting "github.com/bandprotocol/chain/v2/testing"
@@ -16,6 +19,13 @@ func TestExportGenesis(t *testing.T) {
 	s := testutil.NewTestSuite(t)
 	ctx, k := s.Ctx, s.Keeper
 
+	s.MockAccountKeeper.EXPECT().GetModuleAccount(ctx, gomock.Any()).Return(authtypes.AccountI(&authtypes.ModuleAccount{
+		BaseAccount: &authtypes.BaseAccount{Address: "test"},
+	})).AnyTimes()
+	s.MockAccountKeeper.EXPECT().GetModuleAddress(types.ModuleName).Return(sdk.AccAddress{}).AnyTimes()
+	s.MockAccountKeeper.EXPECT().SetModuleAccount(ctx, gomock.Any()).AnyTimes()
+	s.MockBankKeeper.EXPECT().GetAllBalances(ctx, gomock.Any()).Return(sdk.Coins{}).AnyTimes()
+
 	data := types.GenesisState{
 		Params: types.DefaultParams(),
 		Members: []types.Member{

x/bandtss/keeper/keeper.go

@@ -45,6 +45,10 @@ func NewKeeper(
 		panic(fmt.Sprintf("%s module account has not been set", types.ModuleName))
 	}
 
+	if _, err := sdk.AccAddressFromBech32(authority); err != nil {
+		panic(fmt.Errorf("invalid bandtss authority address: %w", err))
+	}
+
 	return &Keeper{
 		cdc:              cdc,
 		storeKey:         storeKey,
@@ -64,6 +68,16 @@ func (k Keeper) GetBandtssAccount(ctx sdk.Context) authtypes.ModuleAccountI {
 	return k.authKeeper.GetModuleAccount(ctx, types.ModuleName)
 }
 
+// GetModuleBalance returns the balance of the bandtss ModuleAccount
+func (k Keeper) GetModuleBalance(ctx sdk.Context) sdk.Coins {
+	return k.bankKeeper.GetAllBalances(ctx, k.GetBandtssAccount(ctx).GetAddress())
+}
+
+// SetModuleAccount sets a module account in the account keeper.
+func (k Keeper) SetModuleAccount(ctx sdk.Context, acc authtypes.ModuleAccountI) {
+	k.authKeeper.SetModuleAccount(ctx, acc)
+}
+
 // Logger gets logger object.
 func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 	return ctx.Logger().With("module", fmt.Sprintf("x/%s", types.ModuleName))

x/bandtss/testutil/mock_expected_keepers.go

@@ -4,7 +4,7 @@ import (
 	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/x/auth/types"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	"github.com/cosmos/cosmos-sdk/x/authz"
 	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
@@ -33,7 +33,8 @@ type AuthzKeeper interface {
 // AccountKeeper defines the expected account keeper (noalias)
 type AccountKeeper interface {
 	GetModuleAddress(name string) sdk.AccAddress
-	GetModuleAccount(ctx sdk.Context, name string) types.ModuleAccountI
+	GetModuleAccount(ctx sdk.Context, name string) authtypes.ModuleAccountI
+	SetModuleAccount(sdk.Context, authtypes.ModuleAccountI)
 }
 
 // BankKeeper defines the expected interface needed to retrieve account balances.

x/bandtss/types/expected_keepers.go

@@ -31,6 +31,10 @@ func NewKeeper(
 	rtr *types.Router,
 	authority string,
 ) *Keeper {
+	if _, err := sdk.AccAddressFromBech32(authority); err != nil {
+		panic(fmt.Errorf("invalid tss authority address: %w", err))
+	}
+
 	return &Keeper{
 		cdc:               cdc,
 		storeKey:          storeKey,