From c1259e7a821f20d7fa9dcbf41f7d97db7e9909b3 Mon Sep 17 00:00:00 2001
From: Dong Lieu <duongduydongtq@gmail.com>
Date: Mon, 15 Jan 2024 15:27:20 +0700
Subject: [PATCH 1/2] fix bug store

---
 x/multi-staking/keeper/store.go      | 10 +++++---
 x/multi-staking/keeper/store_test.go | 38 ++++++++++++++++++++++++++++
 x/multi-staking/keeper/unlock.go     |  2 +-
 x/multi-staking/types/unlock.go      |  2 ++
 4 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/x/multi-staking/keeper/store.go b/x/multi-staking/keeper/store.go
index 47fd6ce4..465a3933 100644
--- 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
index ad843b96..47a811c2 100644
--- 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
index 7ff7be26..128b63b4 100644
--- 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
index 75a4d973..f4668ee9 100644
--- 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),
 		},

From 8fa16fe91f2fc4f44775b38ec6c31fb6fe170ec1 Mon Sep 17 00:00:00 2001
From: Dong Lieu <duongduydongtq@gmail.com>
Date: Mon, 15 Jan 2024 15:34:22 +0700
Subject: [PATCH 2/2] lint

---
 x/multi-staking/keeper/store_test.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/x/multi-staking/keeper/store_test.go b/x/multi-staking/keeper/store_test.go
index 47a811c2..ce9f40aa 100644
--- a/x/multi-staking/keeper/store_test.go
+++ b/x/multi-staking/keeper/store_test.go
@@ -11,7 +11,7 @@ import (
 func (suite *KeeperTestSuite) TestSetBondWeight() {
 	suite.SetupTest()
 
-	gasDenom := "ario"
+	const gasDenom = "ario"
 	govDenom := "arst"
 	gasWeight := sdk.OneDec()
 	govWeight := sdk.NewDecWithPrec(2, 4)
@@ -114,12 +114,12 @@ func (suite *KeeperTestSuite) TestSetAndGetMultiStakingUnlock() {
 		Entries:  Entries,
 	}
 
-	unLocks, found := suite.msKeeper.GetMultiStakingUnlock(suite.ctx, unLockID)
+	_, 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)
+	unLocks, found := suite.msKeeper.GetMultiStakingUnlock(suite.ctx, unLockID)
 	suite.Require().True(found)
 
 	suite.Require().Equal(unLocks.Entries[0].CreationHeight, Entries[0].CreationHeight)