Skip to content

Commit

Permalink
feat: add alias functions and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
RiccardoM committed Jun 18, 2024
1 parent 7893d3e commit cd21618
Show file tree
Hide file tree
Showing 4 changed files with 496 additions and 0 deletions.
76 changes: 76 additions & 0 deletions x/restaking/keeper/alias_functions.go
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
}
282 changes: 282 additions & 0 deletions x/restaking/keeper/alias_functions_test.go
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)
}
})
}
}
Loading

0 comments on commit cd21618

Please sign in to comment.