refactor: declare global const for rewards ratio and add tests

x/distribution/keeper/common_test.go

@@ -24,6 +24,7 @@ import (
 	assetskeeper "github.com/milkyway-labs/milkyway/v7/x/assets/keeper"
 	"github.com/milkyway-labs/milkyway/v7/x/distribution/keeper"
 	"github.com/milkyway-labs/milkyway/v7/x/distribution/testutils"
+	"github.com/milkyway-labs/milkyway/v7/x/distribution/types"
 	operatorskeeper "github.com/milkyway-labs/milkyway/v7/x/operators/keeper"
 	poolskeeper "github.com/milkyway-labs/milkyway/v7/x/pools/keeper"
 	serviceskeeper "github.com/milkyway-labs/milkyway/v7/x/services/keeper"
@@ -73,6 +74,9 @@ func (suite *KeeperTestSuite) SetupTest() {
 
 	err := suite.stakingKeeper.SetParams(suite.ctx, stakingtypes.DefaultParams())
 	suite.Require().NoError(err)
+
+	// Reset to the default(50%) for every test
+	types.VestingAccountRewardsRatio = sdkmath.LegacyNewDecWithPrec(5, 1)
 }
 
 // fundAccount adds the given amount of coins to the account with the given address

x/distribution/keeper/hooks.go

@@ -8,6 +8,8 @@ import (
 	vestingexported "github.com/cosmos/cosmos-sdk/x/auth/vesting/exported"
 	distrkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+
+	"github.com/milkyway-labs/milkyway/v7/x/distribution/types"
 )
 
 var _ stakingtypes.StakingHooks = Hooks{}
@@ -58,7 +60,7 @@ func (h Hooks) BeforeDelegationSharesModified(ctx context.Context, delAddr sdk.A
 	acc := h.k.accountKeeper.GetAccount(ctx, delAddr)
 	_, isVestingAcc := acc.(vestingexported.VestingAccount)
 	if isVestingAcc {
-		shares = shares.QuoInt64(2) // 50%
+		shares = shares.MulTruncate(types.VestingAccountRewardsRatio)
 	}
 
 	validator, err := h.k.Validators.Get(ctx, valAddr)
@@ -85,7 +87,7 @@ func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddre
 	acc := h.k.accountKeeper.GetAccount(ctx, delAddr)
 	_, isVestingAcc := acc.(vestingexported.VestingAccount)
 	if isVestingAcc {
-		shares = shares.QuoInt64(2)
+		shares = shares.MulTruncate(types.VestingAccountRewardsRatio)
 	}
 
 	validator, err := h.k.Validators.Get(ctx, valAddr)

x/distribution/keeper/keeper.go

@@ -21,7 +21,7 @@ type Keeper struct {
 	accountKeeper distrtypes.AccountKeeper
 	stakingKeeper types.StakingKeeper
 
-	Validators    collections.Map[sdk.ValAddress, stakingtypes.Validator]
+	Validators collections.Map[sdk.ValAddress, stakingtypes.Validator]
 }
 
 func NewKeeper(
@@ -72,7 +72,7 @@ func (sk StakingKeeper) Delegation(ctx context.Context, delAddr sdk.AccAddress,
 	acc := sk.k.accountKeeper.GetAccount(ctx, delAddr)
 	_, isVestingAcc := acc.(vestingexported.VestingAccount)
 	if isVestingAcc {
-		delegation.Shares = delegation.Shares.QuoInt64(2)
+		delegation.Shares = delegation.Shares.MulTruncate(types.VestingAccountRewardsRatio)
 	}
 	return delegation, nil
 }

x/distribution/keeper/rewards_test.go

@@ -3,71 +3,103 @@ package keeper_test
 import (
 	"time"
 
+	sdkmath "cosmossdk.io/math"
 	distrkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
 	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	"github.com/milkyway-labs/milkyway/v7/app/testutil"
 	"github.com/milkyway-labs/milkyway/v7/utils"
+	"github.com/milkyway-labs/milkyway/v7/x/distribution/types"
 )
 
 func (suite *KeeperTestSuite) TestVestingAccountRewards() {
-	ctx, _ := suite.ctx.CacheContext()
-	ctx = ctx.WithBlockTime(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC))
+	testCases := []struct {
+		name                     string
+		rewardsRatio             sdkmath.LegacyDec
+		expNormalAccountRewards  string
+		expVestingAccountRewards string
+	}{
+		{
+			name:                     "0% vesting account rewards ratio",
+			rewardsRatio:             sdkmath.LegacyZeroDec(),
+			expNormalAccountRewards:  "500000.000000000000000000stake",
+			expVestingAccountRewards: "",
+		},
+		{
+			name:                     "10% vesting account rewards ratio",
+			rewardsRatio:             utils.MustParseDec("0.1"),
+			expNormalAccountRewards:  "476190.476190476190000000stake", // 1000000 * 1 / 2.1
+			expVestingAccountRewards: "47619.047619047619000000stake",  // 1000000 * 0.1 / 2.1
+		},
+		{
+			name:                     "50% vesting account rewards ratio",
+			rewardsRatio:             utils.MustParseDec("0.5"),
+			expNormalAccountRewards:  "400000.000000000000000000stake",
+			expVestingAccountRewards: "200000.000000000000000000stake",
+		},
+		{
+			name:                     "100% vesting account rewards ratio",
+			rewardsRatio:             sdkmath.LegacyOneDec(),
+			expNormalAccountRewards:  "333333.333333333333000000stake",
+			expVestingAccountRewards: "333333.333333333333000000stake",
+		},
+	}
 
-	valAddr := testutil.TestAddress(10000)
-	validator := suite.createValidator(
-		ctx,
-		valAddr,
-		stakingtypes.NewCommissionRates(utils.MustParseDec("0"), utils.MustParseDec("0.2"), utils.MustParseDec("0.01")),
-		utils.MustParseCoin("1000000stake"),
-		true,
-	)
+	for _, tc := range testCases {
+		suite.Run(tc.name, func() {
+			suite.SetupTest()
+			// Override the rewards ratio
+			types.VestingAccountRewardsRatio = tc.rewardsRatio
 
-	delAddr1 := testutil.TestAddress(1)
-	suite.delegate(ctx, delAddr1.String(), validator.GetOperator(), utils.MustParseCoin("1000000stake"), true)
-	delAddr2 := testutil.TestAddress(2)
-	vestingEndTime := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
-	suite.createVestingAccount(
-		ctx,
-		testutil.TestAddress(10001).String(),
-		delAddr2.String(),
-		utils.MustParseCoins("1000000stake"),
-		vestingEndTime.Unix(),
-		false,
-		true,
-	)
-	suite.delegate(ctx, delAddr2.String(), validator.GetOperator(), utils.MustParseCoin("1000000stake"), false)
+			ctx := suite.ctx.WithBlockTime(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC))
 
-	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1).WithBlockTime(ctx.BlockTime().Add(5 * time.Second))
-	err := suite.k.AllocateTokensToValidator(ctx, validator, utils.MustParseDecCoins("1000000stake"))
-	suite.Require().NoError(err)
+			normalAddr := testutil.TestAddress(1)
+			vestingAddr := testutil.TestAddress(2)
+			vestingEndTime := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+			suite.createVestingAccount(
+				ctx,
+				testutil.TestAddress(10001).String(),
+				vestingAddr.String(),
+				utils.MustParseCoins("1000000stake"),
+				vestingEndTime.Unix(),
+				false,
+				true,
+			)
 
-	querier := distrkeeper.NewQuerier(suite.k.Keeper)
-	// Delegator 1 receive 40% of rewards
-	cacheCtx, _ := ctx.CacheContext()
-	rewards, err := querier.DelegationRewards(cacheCtx, &distrtypes.QueryDelegationRewardsRequest{
-		DelegatorAddress: delAddr1.String(),
-		ValidatorAddress: validator.GetOperator(),
-	})
-	suite.Require().NoError(err)
-	suite.Require().Equal("400000.000000000000000000stake", rewards.Rewards.String())
+			valAddr := testutil.TestAddress(10000)
+			validator := suite.createValidator(
+				ctx,
+				valAddr,
+				stakingtypes.NewCommissionRates(utils.MustParseDec("0"), utils.MustParseDec("0.2"), utils.MustParseDec("0.01")),
+				utils.MustParseCoin("1000000stake"),
+				true,
+			)
+			suite.delegate(ctx, normalAddr.String(), validator.GetOperator(), utils.MustParseCoin("1000000stake"), true)
+			suite.delegate(ctx, vestingAddr.String(), validator.GetOperator(), utils.MustParseCoin("1000000stake"), false)
 
-	// Delegator 2 has a vesting account so it receives 20% of rewards
-	cacheCtx, _ = ctx.CacheContext()
-	rewards, err = querier.DelegationRewards(cacheCtx, &distrtypes.QueryDelegationRewardsRequest{
-		DelegatorAddress: delAddr2.String(),
-		ValidatorAddress: validator.GetOperator(),
-	})
-	suite.Require().NoError(err)
-	suite.Require().Equal("200000.000000000000000000stake", rewards.Rewards.String())
+			ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1).WithBlockTime(ctx.BlockTime().Add(5 * time.Second))
+			err := suite.k.AllocateTokensToValidator(ctx, validator, utils.MustParseDecCoins("1000000stake"))
+			suite.Require().NoError(err)
 
-	// The rest of the rewards was allocated to the validator
-	cacheCtx, _ = ctx.CacheContext()
-	rewards, err = querier.DelegationRewards(cacheCtx, &distrtypes.QueryDelegationRewardsRequest{
-		DelegatorAddress: valAddr.String(),
-		ValidatorAddress: validator.GetOperator(),
-	})
-	suite.Require().NoError(err)
-	suite.Require().Equal("400000.000000000000000000stake", rewards.Rewards.String())
+			querier := distrkeeper.NewQuerier(suite.k.Keeper)
+			// Query the normal account's rewards
+			cacheCtx, _ := ctx.CacheContext()
+			rewards, err := querier.DelegationRewards(cacheCtx, &distrtypes.QueryDelegationRewardsRequest{
+				DelegatorAddress: normalAddr.String(),
+				ValidatorAddress: validator.GetOperator(),
+			})
+			suite.Require().NoError(err)
+			suite.Require().Equal(tc.expNormalAccountRewards, rewards.Rewards.String())
+
+			// Query the vesting account's rewards
+			cacheCtx, _ = ctx.CacheContext()
+			rewards, err = querier.DelegationRewards(cacheCtx, &distrtypes.QueryDelegationRewardsRequest{
+				DelegatorAddress: vestingAddr.String(),
+				ValidatorAddress: validator.GetOperator(),
+			})
+			suite.Require().NoError(err)
+			suite.Require().Equal(tc.expVestingAccountRewards, rewards.Rewards.String())
+		})
+	}
 }

x/distribution/types/rewards.go

@@ -0,0 +1,9 @@
+package types
+
+import (
+	sdkmath "cosmossdk.io/math"
+)
+
+var (
+	VestingAccountRewardsRatio = sdkmath.LegacyNewDecWithPrec(5, 1) // 50%
+)