fix bug store

realiotech · Jan 15, 2024 · c1259e7 · c1259e7
1 parent d96beeb
commit c1259e7
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 4 deletions.
diff --git a/x/multi-staking/keeper/store.go b/x/multi-staking/keeper/store.go
@@ -155,8 +155,8 @@ func (k Keeper) BondWeightIterator(ctx sdk.Context, cb func(denom string, bondWe
 
 func (k Keeper) GetMultiStakingUnlock(ctx sdk.Context, multiStakingUnlockID types.UnlockID) (unlock types.MultiStakingUnlock, found bool) {
 	store := ctx.KVStore(k.storeKey)
-	value := store.Get(multiStakingUnlockID.ToBytes())
-
+	key := append(multiStakingUnlockID.ToBytes(), []byte{0x1}...)
+	value := store.Get(key)
 	if value == nil {
 		return unlock, false
 	}
@@ -172,8 +172,12 @@ func (k Keeper) SetMultiStakingUnlock(ctx sdk.Context, unlock types.MultiStaking
 	store := ctx.KVStore(k.storeKey)
 
 	bz := k.cdc.MustMarshal(&unlock)
+	if unlock.UnlockID == nil {
+		panic("unlock.UnlockID cannot be nil")
+	}
+	key := append(unlock.UnlockID.ToBytes(), []byte{0x1}...)
 
-	store.Set(unlock.UnlockID.ToBytes(), bz)
+	store.Set(key, bz)
 }
 
 func (k Keeper) DeleteMultiStakingUnlock(ctx sdk.Context, unlockID types.UnlockID) {

diff --git a/x/multi-staking/keeper/store_test.go b/x/multi-staking/keeper/store_test.go
@@ -3,6 +3,7 @@ package keeper_test
 import (
 	"github.com/realio-tech/multi-staking-module/testutil"
 	multistakingkeeper "github.com/realio-tech/multi-staking-module/x/multi-staking/keeper"
+	mulStakingtypes "github.com/realio-tech/multi-staking-module/x/multi-staking/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -86,3 +87,40 @@ func (suite *KeeperTestSuite) TestSetValidatorMultiStakingCoin() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestSetAndGetMultiStakingUnlock() {
+	suite.SetupTest()
+	val := testutil.GenValAddress()
+	del := testutil.GenAddress()
+	denom := "ario"
+
+	// set:
+	unLockID := mulStakingtypes.UnlockID{
+		MultiStakerAddr: del.String(),
+		ValAddr:         val.String(),
+	}
+
+	Entries := []mulStakingtypes.UnlockEntry{
+		{
+			CreationHeight: 1,
+			UnlockingCoin: mulStakingtypes.MultiStakingCoin{
+				Denom: denom,
+			},
+		},
+	}
+
+	mulStakingUnllock := mulStakingtypes.MultiStakingUnlock{
+		UnlockID: &unLockID,
+		Entries:  Entries,
+	}
+
+	unLocks, found := suite.msKeeper.GetMultiStakingUnlock(suite.ctx, unLockID)
+	suite.Require().False(found)
+
+	suite.msKeeper.SetMultiStakingUnlock(suite.ctx, mulStakingUnllock)
+
+	unLocks, found = suite.msKeeper.GetMultiStakingUnlock(suite.ctx, unLockID)
+	suite.Require().True(found)
+
+	suite.Require().Equal(unLocks.Entries[0].CreationHeight, Entries[0].CreationHeight)
+}
diff --git a/x/multi-staking/keeper/unlock.go b/x/multi-staking/keeper/unlock.go
@@ -43,7 +43,7 @@ func (k Keeper) SetMultiStakingUnlockEntry(
 	if found {
 		unlock.AddEntry(ctx.BlockHeight(), multistakingCoin)
 	} else {
-		unlock = types.NewMultiStakingUnlock(ctx.BlockHeight(), multistakingCoin)
+		unlock = types.NewMultiStakingUnlock(&unlockID, ctx.BlockHeight(), multistakingCoin)
 	}
 
 	k.SetMultiStakingUnlock(ctx, unlock)

diff --git a/x/multi-staking/types/unlock.go b/x/multi-staking/types/unlock.go
@@ -27,9 +27,11 @@ func (e UnlockEntry) String() string {
 //
 //nolint:interfacer
 func NewMultiStakingUnlock(
+	unlockID *UnlockID,
 	creationHeight int64, weightedCoin MultiStakingCoin,
 ) MultiStakingUnlock {
 	return MultiStakingUnlock{
+		UnlockID: unlockID,
 		Entries: []UnlockEntry{
 			NewUnlockEntry(creationHeight, weightedCoin),
 		},