From 75dd4305e35ce23663694de8a89636ac33ed3b82 Mon Sep 17 00:00:00 2001
From: Dzung Do | Decentrio <neitdung@gmail.com>
Date: Tue, 16 Jan 2024 21:11:43 +0700
Subject: [PATCH] !fix:  fix LockID.ToBytes and UnlockID.ToBytes (#108)

* Fix key prefix and add key tests

* fix lint
---
 x/multi-staking/keeper/store.go   |  6 ++---
 x/multi-staking/types/key.go      | 24 +++++++++----------
 x/multi-staking/types/key_test.go | 40 +++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 15 deletions(-)
 create mode 100644 x/multi-staking/types/key_test.go

diff --git a/x/multi-staking/keeper/store.go b/x/multi-staking/keeper/store.go
index 47fd6ce4..808800b4 100644
--- a/x/multi-staking/keeper/store.go
+++ b/x/multi-staking/keeper/store.go
@@ -75,7 +75,7 @@ func (k Keeper) ValidatorMultiStakingCoinIterator(ctx sdk.Context, cb func(valAd
 func (k Keeper) GetMultiStakingLock(ctx sdk.Context, multiStakingLockID types.LockID) (types.MultiStakingLock, bool) {
 	store := ctx.KVStore(k.storeKey)
 
-	bz := store.Get(multiStakingLockID.ToByte())
+	bz := store.Get(multiStakingLockID.ToBytes())
 	if bz == nil {
 		return types.MultiStakingLock{}, false
 	}
@@ -95,13 +95,13 @@ func (k Keeper) SetMultiStakingLock(ctx sdk.Context, multiStakingLock types.Mult
 
 	bz := k.cdc.MustMarshal(&multiStakingLock)
 
-	store.Set(multiStakingLock.LockID.ToByte(), bz)
+	store.Set(multiStakingLock.LockID.ToBytes(), bz)
 }
 
 func (k Keeper) RemoveMultiStakingLock(ctx sdk.Context, multiStakingLockID types.LockID) {
 	store := ctx.KVStore(k.storeKey)
 
-	store.Delete(multiStakingLockID.ToByte())
+	store.Delete(multiStakingLockID.ToBytes())
 }
 
 func (k Keeper) MultiStakingLockIterator(ctx sdk.Context, cb func(stakingLock types.MultiStakingLock) (stop bool)) {
diff --git a/x/multi-staking/types/key.go b/x/multi-staking/types/key.go
index 3d21ec53..e737e589 100644
--- a/x/multi-staking/types/key.go
+++ b/x/multi-staking/types/key.go
@@ -55,21 +55,21 @@ func MultiStakingUnlockID(multiStakerAddr string, valAddr string) UnlockID {
 }
 
 func DelAddrAndValAddrFromLockID(lockID []byte) (multiStakerAddr sdk.AccAddress, valAddr sdk.ValAddress) {
-	lenMultiStakerAddr := lockID[0]
+	lenMultiStakerAddr := lockID[1]
 
-	multiStakerAddr = lockID[1 : lenMultiStakerAddr+1]
+	multiStakerAddr = lockID[2 : lenMultiStakerAddr+2]
 
-	valAddr = lockID[1+lenMultiStakerAddr:]
+	valAddr = lockID[2+lenMultiStakerAddr:]
 
 	return multiStakerAddr, valAddr
 }
 
 func DelAddrAndValAddrFromUnlockID(unlockID []byte) (multiStakerAddr sdk.AccAddress, valAddr sdk.ValAddress) {
-	lenMultiStakerAddr := unlockID[0]
+	lenMultiStakerAddr := unlockID[1]
 
-	multiStakerAddr = unlockID[1 : lenMultiStakerAddr+1]
+	multiStakerAddr = unlockID[2 : lenMultiStakerAddr+2]
 
-	valAddr = unlockID[1+lenMultiStakerAddr:]
+	valAddr = unlockID[2+lenMultiStakerAddr:]
 
 	return multiStakerAddr, valAddr
 }
@@ -80,7 +80,7 @@ func DelAddrAndValAddrFromUnlockID(unlockID []byte) (multiStakerAddr sdk.AccAddr
 // 	return append(GetUBDsKey(delAddr.Bytes()), address.MustLengthPrefix(valAddr)...)
 // }
 
-func (l LockID) ToByte() []byte {
+func (l LockID) ToBytes() []byte {
 	multiStakerAddr, valAcc, err := AccAddrAndValAddrFromStrings(l.MultiStakerAddr, l.ValAddr)
 	if err != nil {
 		panic(err)
@@ -92,9 +92,9 @@ func (l LockID) ToByte() []byte {
 
 	DVPair[0] = uint8(lenMultiStakerAddr)
 
-	copy(multiStakerAddr[:], DVPair[1:])
+	copy(DVPair[1:], multiStakerAddr[:])
 
-	copy(multiStakerAddr[:], DVPair[1+lenMultiStakerAddr:])
+	copy(DVPair[1+lenMultiStakerAddr:], valAcc[:])
 
 	return append(MultiStakingLockPrefix, DVPair...)
 }
@@ -111,9 +111,9 @@ func (l UnlockID) ToBytes() []byte {
 
 	DVPair[0] = uint8(lenMultiStakerAddr)
 
-	copy(multiStakerAddr[:], DVPair[1:])
+	copy(DVPair[1:], multiStakerAddr[:])
 
-	copy(multiStakerAddr[:], DVPair[1+lenMultiStakerAddr:])
+	copy(DVPair[1+lenMultiStakerAddr:], valAcc[:])
 
-	return append(MultiStakingLockPrefix, DVPair...)
+	return append(MultiStakingUnlockPrefix, DVPair...)
 }
diff --git a/x/multi-staking/types/key_test.go b/x/multi-staking/types/key_test.go
new file mode 100644
index 00000000..6e54d33f
--- /dev/null
+++ b/x/multi-staking/types/key_test.go
@@ -0,0 +1,40 @@
+package types_test
+
+import (
+	"testing"
+
+	"github.com/realio-tech/multi-staking-module/testutil"
+	multistakingtypes "github.com/realio-tech/multi-staking-module/x/multi-staking/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDelAddrAndValAddrFromLockID(t *testing.T) {
+	val := testutil.GenValAddress()
+	del := testutil.GenAddress()
+
+	lockID := multistakingtypes.MultiStakingLockID(del.String(), val.String())
+	lockBytes := lockID.ToBytes()
+	rsDel, rsVal := multistakingtypes.DelAddrAndValAddrFromLockID(lockBytes)
+
+	require.Equal(t, del, rsDel)
+	require.Equal(t, val, rsVal)
+}
+
+func TestMultiStakingLockIterator(t *testing.T) {
+	val := testutil.GenValAddress()
+	delA := testutil.GenAddress()
+	delB := testutil.GenAddress()
+
+	lockIDA := multistakingtypes.LockID{
+		MultiStakerAddr: delA.String(),
+		ValAddr:         val.String(),
+	}
+
+	lockIDB := multistakingtypes.LockID{
+		MultiStakerAddr: delB.String(),
+		ValAddr:         val.String(),
+	}
+
+	require.NotEqual(t, lockIDA, lockIDB)
+	require.NotEqual(t, lockIDA.ToBytes(), lockIDB.ToBytes())
+}