-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
4 changed files
with
496 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,76 @@ | ||
package keeper | ||
|
||
import ( | ||
storetypes "cosmossdk.io/store/types" | ||
sdk "github.com/cosmos/cosmos-sdk/types" | ||
|
||
"github.com/milkyway-labs/milkyway/x/restaking/types" | ||
) | ||
|
||
// IterateAllPoolDelegations iterates over all the pool delegations and performs a callback function | ||
func (k *Keeper) IterateAllPoolDelegations(ctx sdk.Context, cb func(delegation types.PoolDelegation) (stop bool)) { | ||
store := ctx.KVStore(k.storeKey) | ||
iterator := store.Iterator(types.PoolDelegationPrefix, storetypes.PrefixEndBytes(types.PoolDelegationPrefix)) | ||
defer iterator.Close() | ||
|
||
for ; iterator.Valid(); iterator.Next() { | ||
delegation := types.MustUnmarshalPoolDelegation(k.cdc, iterator.Value()) | ||
if cb(delegation) { | ||
break | ||
} | ||
} | ||
} | ||
|
||
// GetAllPoolDelegations returns all the pool delegations | ||
func (k *Keeper) GetAllPoolDelegations(ctx sdk.Context) []types.PoolDelegation { | ||
var delegations []types.PoolDelegation | ||
k.IterateAllPoolDelegations(ctx, func(delegation types.PoolDelegation) bool { | ||
delegations = append(delegations, delegation) | ||
return false | ||
}) | ||
return delegations | ||
} | ||
|
||
// GetAllDelegatorPoolDelegations returns all the pool delegations of a given delegator | ||
func (k *Keeper) GetAllDelegatorPoolDelegations(ctx sdk.Context, delegator string) []types.PoolDelegation { | ||
store := ctx.KVStore(k.storeKey) | ||
delegatorPrefixKey := types.UserPoolDelegationsStorePrefix(delegator) | ||
|
||
iterator := store.Iterator(delegatorPrefixKey, storetypes.PrefixEndBytes(delegatorPrefixKey)) // Smallest to largest | ||
defer iterator.Close() | ||
|
||
var delegations []types.PoolDelegation | ||
for i := 0; iterator.Valid(); iterator.Next() { | ||
delegation := types.MustUnmarshalPoolDelegation(k.cdc, iterator.Value()) | ||
delegations = append(delegations, delegation) | ||
i++ | ||
} | ||
|
||
return delegations | ||
} | ||
|
||
// GetPoolDelegations returns all the delegations to a given pool | ||
func (k *Keeper) GetPoolDelegations(ctx sdk.Context, poolID uint32) ([]types.PoolDelegation, error) { | ||
store := ctx.KVStore(k.storeKey) | ||
prefix := types.DelegationsByPoolIDStorePrefix(poolID) | ||
iterator := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix)) | ||
defer iterator.Close() | ||
|
||
var delegations []types.PoolDelegation | ||
for ; iterator.Valid(); iterator.Next() { | ||
_, delegatorAddress, err := types.ParseDelegationsByPoolIDKey(iterator.Key()) | ||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
bz := store.Get(types.UserPoolDelegationStoreKey(delegatorAddress, poolID)) | ||
delegation, err := types.UnmarshalPoolDelegation(k.cdc, bz) | ||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
delegations = append(delegations, delegation) | ||
} | ||
|
||
return delegations, nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,282 @@ | ||
package keeper_test | ||
|
||
import ( | ||
sdkmath "cosmossdk.io/math" | ||
sdk "github.com/cosmos/cosmos-sdk/types" | ||
|
||
"github.com/milkyway-labs/milkyway/x/restaking/types" | ||
) | ||
|
||
func (suite *KeeperTestSuite) TestKeeper_GetAllPoolDelegations() { | ||
testCases := []struct { | ||
name string | ||
setup func() | ||
store func(ctx sdk.Context) | ||
shouldErr bool | ||
expDelegations []types.PoolDelegation | ||
check func(ctx sdk.Context) | ||
}{ | ||
{ | ||
name: "delegations are returned properly", | ||
store: func(ctx sdk.Context) { | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
}, | ||
expDelegations: []types.PoolDelegation{ | ||
types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
), | ||
types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
), | ||
types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
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) | ||
} | ||
|
||
delegations := suite.k.GetAllPoolDelegations(ctx) | ||
suite.Require().Equal(tc.expDelegations, delegations) | ||
|
||
if tc.check != nil { | ||
tc.check(ctx) | ||
} | ||
}) | ||
} | ||
} | ||
|
||
func (suite *KeeperTestSuite) TestKeeper_GetAllDelegatorPoolDelegations() { | ||
testCases := []struct { | ||
name string | ||
setup func() | ||
store func(ctx sdk.Context) | ||
delegator string | ||
expDelegations []types.PoolDelegation | ||
check func(ctx sdk.Context) | ||
}{ | ||
{ | ||
name: "user without delegations returns empty list", | ||
delegator: "cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
expDelegations: nil, | ||
}, | ||
{ | ||
name: "user with single delegation returns it properly", | ||
store: func(ctx sdk.Context) { | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
}, | ||
delegator: "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
expDelegations: []types.PoolDelegation{ | ||
types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
sdkmath.LegacyNewDec(100), | ||
), | ||
}, | ||
}, | ||
{ | ||
name: "user with multiple delegations returns them properly", | ||
store: func(ctx sdk.Context) { | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
}, | ||
delegator: "cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
expDelegations: []types.PoolDelegation{ | ||
types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
), | ||
types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
), | ||
}, | ||
}, | ||
} | ||
|
||
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) | ||
} | ||
|
||
delegations := suite.k.GetAllDelegatorPoolDelegations(ctx, tc.delegator) | ||
suite.Require().Equal(tc.expDelegations, delegations) | ||
|
||
if tc.check != nil { | ||
tc.check(ctx) | ||
} | ||
}) | ||
} | ||
} | ||
|
||
func (suite *KeeperTestSuite) TestKeeper_GetPoolDelegations() { | ||
testCases := []struct { | ||
name string | ||
setup func() | ||
store func(ctx sdk.Context) | ||
poolID uint32 | ||
shouldErr bool | ||
expDelegations []types.PoolDelegation | ||
check func(ctx sdk.Context) | ||
}{ | ||
{ | ||
name: "pool without delegations returns empty list", | ||
poolID: 1, | ||
shouldErr: false, | ||
expDelegations: nil, | ||
}, | ||
{ | ||
name: "pool with single delegation returns it properly", | ||
store: func(ctx sdk.Context) { | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
}, | ||
poolID: 2, | ||
shouldErr: false, | ||
expDelegations: []types.PoolDelegation{ | ||
types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
), | ||
}, | ||
}, | ||
{ | ||
name: "pool with multiple delegations returns them properly", | ||
store: func(ctx sdk.Context) { | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
2, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(50), | ||
)) | ||
suite.k.SavePoolDelegation(ctx, types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
sdkmath.LegacyNewDec(100), | ||
)) | ||
}, | ||
poolID: 1, | ||
shouldErr: false, | ||
expDelegations: []types.PoolDelegation{ | ||
types.NewPoolDelegation( | ||
1, | ||
"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd", | ||
sdkmath.LegacyNewDec(100), | ||
), | ||
types.NewPoolDelegation( | ||
1, | ||
"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4", | ||
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) | ||
} | ||
|
||
delegations, err := suite.k.GetPoolDelegations(ctx, tc.poolID) | ||
if tc.shouldErr { | ||
suite.Require().Error(err) | ||
} else { | ||
suite.Require().NoError(err) | ||
suite.Require().Equal(tc.expDelegations, delegations) | ||
} | ||
|
||
if tc.check != nil { | ||
tc.check(ctx) | ||
} | ||
}) | ||
} | ||
} |
Oops, something went wrong.