From c41706021f293416cb8981b368e1c492abf71100 Mon Sep 17 00:00:00 2001 From: Riccardo Montagnin Date: Tue, 18 Jun 2024 11:22:00 -0500 Subject: [PATCH] feat(tests): add pool restaking tests --- x/restaking/keeper/common_test.go | 46 ++- x/restaking/keeper/pool_restaking.go | 2 +- x/restaking/keeper/pool_restaking_test.go | 475 ++++++++++++++++++++++ x/restaking/types/keys.go | 6 +- 4 files changed, 524 insertions(+), 5 deletions(-) create mode 100644 x/restaking/keeper/pool_restaking_test.go diff --git a/x/restaking/keeper/common_test.go b/x/restaking/keeper/common_test.go index 518d3832..b0cbcd9e 100644 --- a/x/restaking/keeper/common_test.go +++ b/x/restaking/keeper/common_test.go @@ -102,5 +102,49 @@ func (suite *KeeperTestSuite) SetupTest() { suite.bk, suite.pk, authorityAddr, - ) + ).SetHooks(newMockHooks()) +} + +// -------------------------------------------------------------------------------------------------------------------- + +// fundAccount adds the given amount of coins to the account with the given address +func (suite *KeeperTestSuite) fundAccount(ctx sdk.Context, address string, amount sdk.Coins) { + // Mint the coins + moduleAcc := suite.ak.GetModuleAccount(ctx, authtypes.Minter) + + err := suite.bk.MintCoins(ctx, moduleAcc.GetName(), amount) + suite.Require().NoError(err) + + // Get the amount to the user + userAddress, err := sdk.AccAddressFromBech32(address) + suite.Require().NoError(err) + err = suite.bk.SendCoinsFromModuleToAccount(ctx, moduleAcc.GetName(), userAddress, amount) + suite.Require().NoError(err) +} + +// -------------------------------------------------------------------------------------------------------------------- + +var _ types.RestakingHooks = &mockHooks{} + +type mockHooks struct { + CalledMap map[string]bool +} + +func newMockHooks() *mockHooks { + return &mockHooks{CalledMap: make(map[string]bool)} +} + +func (m mockHooks) BeforePoolDelegationCreated(ctx sdk.Context, poolID uint32, delegator string) error { + m.CalledMap["BeforePoolDelegationCreated"] = true + return nil +} + +func (m mockHooks) BeforePoolDelegationSharesModified(ctx sdk.Context, poolID uint32, delegator string) error { + m.CalledMap["BeforePoolDelegationSharesModified"] = true + return nil +} + +func (m mockHooks) AfterPoolDelegationModified(ctx sdk.Context, poolID uint32, delegator string) error { + m.CalledMap["AfterPoolDelegationModified"] = true + return nil } diff --git a/x/restaking/keeper/pool_restaking.go b/x/restaking/keeper/pool_restaking.go index 01d14295..558913be 100644 --- a/x/restaking/keeper/pool_restaking.go +++ b/x/restaking/keeper/pool_restaking.go @@ -17,7 +17,7 @@ func (k *Keeper) SavePoolDelegation(ctx sdk.Context, delegation types.PoolDelega store.Set(types.UserPoolDelegationStoreKey(delegation.UserAddress, delegation.PoolID), delegationBz) // Store the delegation in the delegations by pool ID store - store.Set(types.DelegationsByPoolIDStoreKey(delegation.PoolID, delegation.UserAddress), []byte{}) + store.Set(types.DelegationByPoolIDStoreKey(delegation.PoolID, delegation.UserAddress), []byte{}) } // GetPoolDelegation retrieves the delegation for the given user and pool diff --git a/x/restaking/keeper/pool_restaking_test.go b/x/restaking/keeper/pool_restaking_test.go new file mode 100644 index 00000000..622ba4b7 --- /dev/null +++ b/x/restaking/keeper/pool_restaking_test.go @@ -0,0 +1,475 @@ +package keeper_test + +import ( + sdkmath "cosmossdk.io/math" + sdk "github.com/cosmos/cosmos-sdk/types" + + poolstypes "github.com/milkyway-labs/milkyway/x/pools/types" + "github.com/milkyway-labs/milkyway/x/restaking/types" +) + +func (suite *KeeperTestSuite) TestKeeper_SavePoolDelegation() { + testCases := []struct { + name string + setup func() + store func(ctx sdk.Context) + delegation types.PoolDelegation + check func(ctx sdk.Context) + }{ + { + name: "pool delegation is stored properly", + delegation: types.NewPoolDelegation( + 1, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdkmath.LegacyNewDec(100), + ), + check: func(ctx sdk.Context) { + store := ctx.KVStore(suite.storeKey) + + // Make sure the user-pool delegation key exists and contains the delegation + delegationBz := store.Get(types.UserPoolDelegationStoreKey("cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", 1)) + suite.Require().NotNil(delegationBz) + + delegation, err := types.UnmarshalPoolDelegation(suite.cdc, delegationBz) + suite.Require().NoError(err) + + suite.Require().Equal(types.NewPoolDelegation( + 1, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdkmath.LegacyNewDec(100), + ), delegation) + + // Make sure the pool-user delegation key exists + hasDelegationsByPoolKey := store.Has(types.DelegationByPoolIDStoreKey(1, "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4")) + suite.Require().True(hasDelegationsByPoolKey) + }, + }, + } + + for _, tc := range testCases { + tc := tc + suite.Run(tc.name, func() { + ctx, _ := suite.ctx.CacheContext() + if tc.setup != nil { + tc.setup() + } + if tc.store != nil { + tc.store(ctx) + } + + suite.k.SavePoolDelegation(ctx, tc.delegation) + + if tc.check != nil { + tc.check(ctx) + } + }) + } +} + +func (suite *KeeperTestSuite) TestKeeper_GetPoolDelegation() { + testCases := []struct { + name string + setup func() + store func(ctx sdk.Context) + poolID uint32 + userAddress string + expFound bool + expDelegation types.PoolDelegation + check func(ctx sdk.Context) + }{ + { + name: "not found delegation returns false", + poolID: 1, + userAddress: "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + expFound: false, + }, + { + name: "found delegation is returned properly", + store: func(ctx sdk.Context) { + suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( + 1, + "cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", + sdkmath.LegacyNewDec(100), + )) + }, + poolID: 1, + userAddress: "cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", + expFound: true, + expDelegation: types.NewPoolDelegation( + 1, + "cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", + sdkmath.LegacyNewDec(100), + ), + }, + } + + for _, tc := range testCases { + tc := tc + suite.Run(tc.name, func() { + ctx, _ := suite.ctx.CacheContext() + if tc.setup != nil { + tc.setup() + } + if tc.store != nil { + tc.store(ctx) + } + + delegation, found := suite.k.GetPoolDelegation(ctx, tc.poolID, tc.userAddress) + if !tc.expFound { + suite.Require().False(found) + } else { + suite.Require().True(found) + suite.Require().Equal(tc.expDelegation, delegation) + } + + if tc.check != nil { + tc.check(ctx) + } + }) + } +} + +func (suite *KeeperTestSuite) TestKeeper_AddPoolTokensAndShares() { + testCases := []struct { + name string + setup func() + store func(ctx sdk.Context) + pool poolstypes.Pool + tokensToAdd sdkmath.Int + shouldErr bool + expPool poolstypes.Pool + expAddedShares sdkmath.LegacyDec + check func(ctx sdk.Context) + }{ + { + name: "adding tokens to an empty pool works properly", + pool: poolstypes.NewPool(1, "umilk"), + tokensToAdd: sdkmath.NewInt(100), + shouldErr: false, + expPool: poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(100), + DelegatorShares: sdkmath.LegacyNewDec(100), + }, + expAddedShares: sdkmath.LegacyNewDec(100), + }, + { + name: "adding tokens to a non-empty pool works properly", + pool: poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(50), + DelegatorShares: sdkmath.LegacyNewDec(100), + }, + tokensToAdd: sdkmath.NewInt(20), + shouldErr: false, + expPool: poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(70), + DelegatorShares: sdkmath.LegacyNewDec(140), + }, + expAddedShares: sdkmath.LegacyNewDec(40), + }, + } + + for _, tc := range testCases { + tc := tc + suite.Run(tc.name, func() { + ctx, _ := suite.ctx.CacheContext() + if tc.setup != nil { + tc.setup() + } + if tc.store != nil { + tc.store(ctx) + } + + pool, addedShares, err := suite.k.AddPoolTokensAndShares(ctx, tc.pool, tc.tokensToAdd) + if tc.shouldErr { + suite.Require().Error(err) + } else { + suite.Require().NoError(err) + suite.Require().Equal(tc.expPool, pool) + suite.Require().Equal(tc.expAddedShares, addedShares) + } + + if tc.check != nil { + tc.check(ctx) + } + }) + } +} + +// -------------------------------------------------------------------------------------------------------------------- + +func (suite *KeeperTestSuite) TestKeeper_DelegateToPool() { + testCases := []struct { + name string + setup func() + store func(ctx sdk.Context) + amount sdk.Coin + delegator string + shouldErr bool + expShares sdkmath.LegacyDec + check func(ctx sdk.Context) + }{ + { + name: "invalid exchange rate pool returns error", + store: func(ctx sdk.Context) { + err := suite.pk.SavePool(ctx, poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.ZeroInt(), + DelegatorShares: sdkmath.LegacyNewDec(100), + }) + suite.Require().NoError(err) + }, + amount: sdk.NewCoin("umilk", sdkmath.NewInt(100)), + delegator: "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + shouldErr: true, + }, + { + name: "invalid delegator address returns error", + store: func(ctx sdk.Context) { + err := suite.pk.SavePool(ctx, poolstypes.NewPool(1, "umilk")) + suite.Require().NoError(err) + }, + amount: sdk.NewCoin("umilk", sdkmath.NewInt(100)), + delegator: "invalid", + shouldErr: true, + }, + { + name: "insufficient funds return error", + store: func(ctx sdk.Context) { + // Create the pool + err := suite.pk.SavePool(ctx, poolstypes.NewPool(1, "umilk")) + suite.Require().NoError(err) + + // Set the next pool id + suite.pk.SetNextPoolID(ctx, 2) + + // Send some funds to the user + suite.fundAccount( + ctx, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdk.NewCoins(sdk.NewCoin("umilk", sdkmath.NewInt(50))), + ) + }, + amount: sdk.NewCoin("umilk", sdkmath.NewInt(100)), + delegator: "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + shouldErr: true, + }, + { + name: "delegating to a non-existing pool works properly", + store: func(ctx sdk.Context) { + // Set the next pool id + suite.pk.SetNextPoolID(ctx, 1) + + // Send some funds to the user + suite.fundAccount( + ctx, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdk.NewCoins(sdk.NewCoin("umilk", sdkmath.NewInt(100))), + ) + }, + amount: sdk.NewCoin("umilk", sdkmath.NewInt(100)), + delegator: "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + shouldErr: false, + expShares: sdkmath.LegacyNewDec(100), + check: func(ctx sdk.Context) { + // Make sure the pool now exists + pool, found := suite.pk.GetPool(ctx, 1) + suite.Require().True(found) + suite.Require().Equal(poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(100), + DelegatorShares: sdkmath.LegacyNewDec(100), + }, pool) + + // Make sure the delegation exists + delegation, found := suite.k.GetPoolDelegation(ctx, 1, "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4") + suite.Require().True(found) + suite.Require().Equal(types.NewPoolDelegation( + 1, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdkmath.LegacyNewDec(100), + ), delegation) + + // Make sure the user balance has been reduced properly + userBalance := suite.bk.GetBalance(ctx, sdk.AccAddress("cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4"), "umilk") + suite.Require().Equal(sdk.NewCoin("umilk", sdkmath.NewInt(0)), userBalance) + + // Make sure the pool account balance has increased properly + poolBalance := suite.bk.GetBalance(ctx, poolstypes.GetPoolAddress(1), "umilk") + suite.Require().Equal(sdk.NewCoin("umilk", sdkmath.NewInt(100)), poolBalance) + }, + }, + { + name: "delegating to an existing pool works properly", + store: func(ctx sdk.Context) { + // Create the pool + err := suite.pk.SavePool(ctx, poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(20), + DelegatorShares: sdkmath.LegacyNewDec(100), + }) + suite.Require().NoError(err) + + // Set the correct pool tokens amount + suite.fundAccount( + ctx, + poolstypes.GetPoolAddress(1).String(), + sdk.NewCoins(sdk.NewCoin("umilk", sdkmath.NewInt(20))), + ) + + // Set the next pool id + suite.pk.SetNextPoolID(ctx, 2) + + // Send some funds to the user + suite.fundAccount( + ctx, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdk.NewCoins(sdk.NewCoin("umilk", sdkmath.NewInt(100))), + ) + }, + amount: sdk.NewCoin("umilk", sdkmath.NewInt(100)), + delegator: "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + shouldErr: false, + expShares: sdkmath.LegacyNewDec(500), + check: func(ctx sdk.Context) { + // Make sure the pool now exists + pool, found := suite.pk.GetPool(ctx, 1) + suite.Require().True(found) + suite.Require().Equal(poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(120), + DelegatorShares: sdkmath.LegacyNewDec(600), + }, pool) + + // Make sure the delegation exists + delegation, found := suite.k.GetPoolDelegation(ctx, 1, "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4") + suite.Require().True(found) + suite.Require().Equal(types.NewPoolDelegation( + 1, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdkmath.LegacyNewDec(500), + ), delegation) + + // Make sure the user balance has been reduced properly + userBalance := suite.bk.GetBalance(ctx, sdk.AccAddress("cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4"), "umilk") + suite.Require().Equal(sdk.NewCoin("umilk", sdkmath.NewInt(0)), userBalance) + + // Make sure the pool account balance has increased properly + poolBalance := suite.bk.GetBalance(ctx, poolstypes.GetPoolAddress(1), "umilk") + suite.Require().Equal(sdk.NewCoin("umilk", sdkmath.NewInt(120)), poolBalance) + }, + }, + { + name: "delegating more tokens works properly", + store: func(ctx sdk.Context) { + // Create the pool + err := suite.pk.SavePool(ctx, poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(80), + DelegatorShares: sdkmath.LegacyNewDec(125), + }) + suite.Require().NoError(err) + + // Set the correct pool tokens amount + suite.fundAccount( + ctx, + poolstypes.GetPoolAddress(1).String(), + sdk.NewCoins(sdk.NewCoin("umilk", sdkmath.NewInt(80))), + ) + + // Set the next pool id + suite.pk.SetNextPoolID(ctx, 2) + + // Save the existing delegation + suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( + 1, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdkmath.LegacyNewDec(100), + )) + + // Send some funds to the user + suite.fundAccount( + ctx, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdk.NewCoins(sdk.NewCoin("umilk", sdkmath.NewInt(100))), + ) + }, + amount: sdk.NewCoin("umilk", sdkmath.NewInt(100)), + delegator: "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + shouldErr: false, + expShares: sdkmath.LegacyNewDecWithPrec(15625, 2), + check: func(ctx sdk.Context) { + // Make sure the pool now exists + pool, found := suite.pk.GetPool(ctx, 1) + suite.Require().True(found) + suite.Require().Equal(poolstypes.Pool{ + ID: 1, + Denom: "umilk", + Address: poolstypes.GetPoolAddress(1).String(), + Tokens: sdkmath.NewInt(180), + DelegatorShares: sdkmath.LegacyNewDecWithPrec(28125, 2), + }, pool) + + // Make sure the delegation has been updated properly + delegation, found := suite.k.GetPoolDelegation(ctx, 1, "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4") + suite.Require().True(found) + suite.Require().Equal(types.NewPoolDelegation( + 1, + "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", + sdkmath.LegacyNewDecWithPrec(25625, 2), // 100 (existing) + 156.25 (new) + ), delegation) + + // Make sure the user balance has been reduced properly + userBalance := suite.bk.GetBalance(ctx, sdk.AccAddress("cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4"), "umilk") + suite.Require().Equal(sdk.NewCoin("umilk", sdkmath.NewInt(0)), userBalance) + + // Make sure the pool account balance has increased properly + poolBalance := suite.bk.GetBalance(ctx, poolstypes.GetPoolAddress(1), "umilk") + suite.Require().Equal(sdk.NewCoin("umilk", sdkmath.NewInt(180)), poolBalance) + }, + }, + } + + for _, tc := range testCases { + tc := tc + suite.Run(tc.name, func() { + ctx, _ := suite.ctx.CacheContext() + if tc.setup != nil { + tc.setup() + } + if tc.store != nil { + tc.store(ctx) + } + + shares, err := suite.k.DelegateToPool(ctx, tc.amount, tc.delegator) + if tc.shouldErr { + suite.Require().Error(err) + } else { + suite.Require().NoError(err) + suite.Require().Equal(tc.expShares, shares) + } + + if tc.check != nil { + tc.check(ctx) + } + }) + } +} diff --git a/x/restaking/types/keys.go b/x/restaking/types/keys.go index 4e44fb46..b0527c95 100644 --- a/x/restaking/types/keys.go +++ b/x/restaking/types/keys.go @@ -31,7 +31,7 @@ func UserPoolDelegationsStorePrefix(userAddress string) []byte { return append(PoolDelegationPrefix, []byte(userAddress)...) } -// UserPoolDelegationStoreKey returns the key used to store the delegation of a user to a given pool +// UserPoolDelegationStoreKey returns the key used to store the user -> pool delegation association func UserPoolDelegationStoreKey(delegator string, poolID uint32) []byte { return append(UserPoolDelegationsStorePrefix(delegator), poolstypes.GetPoolIDBytes(poolID)...) } @@ -41,8 +41,8 @@ func DelegationsByPoolIDStorePrefix(poolID uint32) []byte { return append(PoolDelegationsByPoolIDPrefix, poolstypes.GetPoolIDBytes(poolID)...) } -// DelegationsByPoolIDStoreKey returns the key used to store the delegations to a given pool -func DelegationsByPoolIDStoreKey(poolID uint32, delegatorAddress string) []byte { +// DelegationByPoolIDStoreKey returns the key used to store the pool -> user delegation association +func DelegationByPoolIDStoreKey(poolID uint32, delegatorAddress string) []byte { return append(DelegationsByPoolIDStorePrefix(poolID), []byte(delegatorAddress)...) }