add unit test

Author: RogerKSI

x/restake/client/cli/tx.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 
+	sdkmath "cosmossdk.io/math"
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/client/tx"
@@ -121,7 +122,7 @@ func GetTxCmdLockPower() *cobra.Command {
 			}
 
 			key := args[0]
-			amount, ok := sdk.NewIntFromString(args[1])
+			amount, ok := sdkmath.NewIntFromString(args[1])
 			if !ok {
 				return fmt.Errorf("invalid amount")
 			}

x/restake/keeper/grpc_query_test.go

@@ -123,6 +123,10 @@ func (suite *KeeperTestSuite) TestQueryRewards() {
 							Key:     ValidKey2,
 							Rewards: nil,
 						},
+						{
+							Key:     ValidKey3,
+							Rewards: nil,
+						},
 					},
 				}
 			},

x/restake/keeper/keeper_key.go

@@ -1,12 +1,36 @@
 package keeper
 
 import (
+	sdkmath "cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 
 	"github.com/bandprotocol/chain/v2/x/restake/types"
 )
 
+func (k Keeper) GetOrCreateKey(ctx sdk.Context, keyName string) (types.Key, error) {
+	key, err := k.GetKey(ctx, keyName)
+	if err != nil {
+		keyAccAddr, err := k.createKeyAccount(ctx, keyName)
+		if err != nil {
+			return types.Key{}, err
+		}
+
+		key = types.Key{
+			Name:            keyName,
+			PoolAddress:     keyAccAddr.String(),
+			IsActive:        true,
+			TotalPower:      sdkmath.NewInt(0),
+			RewardPerPowers: sdk.NewDecCoins(),
+			Remainders:      sdk.NewDecCoins(),
+		}
+
+		k.SetKey(ctx, key)
+	}
+
+	return key, nil
+}
+
 // AddRewards adds rewards to the pool address and re-calculate reward_per_share of the key
 func (k Keeper) AddRewards(ctx sdk.Context, sender sdk.AccAddress, keyName string, rewards sdk.Coins) error {
 	key, err := k.GetKey(ctx, keyName)
@@ -28,7 +52,7 @@ func (k Keeper) AddRewards(ctx sdk.Context, sender sdk.AccAddress, keyName strin
 	}
 
 	decRewards := sdk.NewDecCoinsFromCoins(rewards.Sort()...)
-	totalPower := sdk.NewDecFromInt(key.TotalPower)
+	totalPower := sdkmath.LegacyNewDecFromInt(key.TotalPower)
 	rewardPerPowers := decRewards.QuoDecTruncate(totalPower)
 	truncatedRewards := decRewards.Sub(rewardPerPowers.MulDecTruncate(totalPower))
 
@@ -50,29 +74,39 @@ func (k Keeper) AddRewards(ctx sdk.Context, sender sdk.AccAddress, keyName strin
 	return nil
 }
 
-func (k Keeper) GetOrCreateKey(ctx sdk.Context, keyName string) (types.Key, error) {
+func (k Keeper) IsActiveKey(ctx sdk.Context, keyName string) bool {
 	key, err := k.GetKey(ctx, keyName)
 	if err != nil {
-		keyAccAddr, err := k.CreateKeyAccount(ctx, keyName)
-		if err != nil {
-			return types.Key{}, err
-		}
+		return false
+	}
 
-		key = types.Key{
-			Name:            keyName,
-			PoolAddress:     keyAccAddr.String(),
-			IsActive:        true,
-			TotalPower:      sdk.NewInt(0),
-			RewardPerPowers: sdk.NewDecCoins(),
-		}
+	return key.IsActive
+}
 
-		k.SetKey(ctx, key)
+func (k Keeper) DeactivateKey(ctx sdk.Context, keyName string) error {
+	key, err := k.GetKey(ctx, keyName)
+	if err != nil {
+		return err
 	}
 
-	return key, nil
+	if !key.IsActive {
+		return types.ErrKeyAlreadyDeactivated
+	}
+
+	key.IsActive = false
+	k.SetKey(ctx, key)
+
+	ctx.EventManager().EmitEvent(
+		sdk.NewEvent(
+			types.EventTypeDeactivateKey,
+			sdk.NewAttribute(types.AttributeKeyKey, keyName),
+		),
+	)
+
+	return nil
 }
 
-func (k Keeper) CreateKeyAccount(ctx sdk.Context, key string) (sdk.AccAddress, error) {
+func (k Keeper) createKeyAccount(ctx sdk.Context, key string) (sdk.AccAddress, error) {
 	header := ctx.BlockHeader()
 
 	buf := []byte(key)
@@ -105,38 +139,6 @@ func (k Keeper) CreateKeyAccount(ctx sdk.Context, key string) (sdk.AccAddress, e
 	return keyAccAddr, nil
 }
 
-func (k Keeper) IsActiveKey(ctx sdk.Context, keyName string) bool {
-	key, err := k.GetKey(ctx, keyName)
-	if err != nil {
-		return false
-	}
-
-	return key.IsActive
-}
-
-func (k Keeper) DeactivateKey(ctx sdk.Context, keyName string) error {
-	key, err := k.GetKey(ctx, keyName)
-	if err != nil {
-		return err
-	}
-
-	if !key.IsActive {
-		return types.ErrKeyAlreadyDeactivated
-	}
-
-	key.IsActive = false
-	k.SetKey(ctx, key)
-
-	ctx.EventManager().EmitEvent(
-		sdk.NewEvent(
-			types.EventTypeDeactivateKey,
-			sdk.NewAttribute(types.AttributeKeyKey, keyName),
-		),
-	)
-
-	return nil
-}
-
 // -------------------------------
 // store part
 // -------------------------------