claim undelegation issue fix

x/multistaking/keeper/delegation.go

@@ -57,6 +57,12 @@ func (k Keeper) SetUndelegation(ctx sdk.Context, undelegation types.Undelegation
 	store.Set(key, k.cdc.MustMarshal(&undelegation))
 }
 
+func (k Keeper) RemoveUndelegation(ctx sdk.Context, id uint64) {
+	store := ctx.KVStore(k.storeKey)
+	key := append(types.KeyPrefixUndelegation, sdk.Uint64ToBigEndian(id)...)
+	store.Delete(key)
+}
+
 func (k Keeper) SetPoolDelegator(ctx sdk.Context, poolId uint64, delegator sdk.AccAddress) {
 	store := ctx.KVStore(k.storeKey)
 	key := append(append(types.KeyPrefixPoolDelegator, sdk.Uint64ToBigEndian(poolId)...), delegator...)

x/multistaking/keeper/msg_server.go

@@ -126,6 +126,8 @@ func (k msgServer) ClaimUndelegation(goCtx context.Context, msg *types.MsgClaimU
 		return nil, err
 	}
 
+	k.keeper.RemoveUndelegation(ctx, undelegation.Id)
+
 	return &types.MsgClaimUndelegationResponse{}, nil
 }
 

x/multistaking/keeper/msg_server_test.go

@@ -0,0 +1,117 @@
+package keeper_test
+
+import (
+	"time"
+
+	multistakingkeeper "github.com/KiraCore/sekai/x/multistaking/keeper"
+	multistakingtypes "github.com/KiraCore/sekai/x/multistaking/types"
+	stakingtypes "github.com/KiraCore/sekai/x/staking/types"
+	"github.com/cometbft/cometbft/crypto/ed25519"
+	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	minttypes "github.com/cosmos/cosmos-sdk/x/mint/types"
+)
+
+// func (k msgServer) ClaimUndelegation(goCtx context.Context, msg *types.MsgClaimUndelegation) (*types.MsgClaimUndelegationResponse, error) {
+// 	ctx := sdk.UnwrapSDKContext(goCtx)
+
+// 	undelegation, found := k.keeper.GetUndelegationById(ctx, msg.UndelegationId)
+// 	if !found {
+// 		return nil, types.ErrUndelegationNotFound
+// 	}
+
+// 	if uint64(ctx.BlockTime().Unix()) < undelegation.Expiry {
+// 		return nil, types.ErrNotEnoughTimePassed
+// 	}
+
+// 	delegator, err := sdk.AccAddressFromBech32(msg.Sender)
+// 	if err != nil {
+// 		return nil, err
+// 	}
+
+// 	err = k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, delegator, undelegation.Amount)
+// 	if err != nil {
+// 		return nil, err
+// 	}
+
+// 	k.keeper.RemoveUndelegation(ctx, undelegation.Id)
+
+// 	return &types.MsgClaimUndelegationResponse{}, nil
+// }
+
+func (suite *KeeperTestSuite) TestClaimUndelegation() {
+	addr1 := sdk.AccAddress(ed25519.GenPrivKey().PubKey().Address().Bytes())
+
+	valAddr := sdk.ValAddress(addr1)
+	pubkeys := simtestutil.CreateTestPubKeys(1)
+	pubKey := pubkeys[0]
+	val, err := stakingtypes.NewValidator(valAddr, pubKey)
+	suite.Require().NoError(err)
+
+	properties := suite.app.CustomGovKeeper.GetNetworkProperties(suite.ctx)
+	properties.MinCollectiveBond = 1
+	suite.app.CustomGovKeeper.SetNetworkProperties(suite.ctx, properties)
+
+	val.Status = stakingtypes.Active
+	suite.app.CustomStakingKeeper.AddValidator(suite.ctx, val)
+
+	stakingPool := multistakingtypes.StakingPool{
+		Id:        1,
+		Validator: valAddr.String(),
+		Enabled:   true,
+	}
+
+	suite.app.MultiStakingKeeper.SetStakingPool(suite.ctx, stakingPool)
+
+	coins := sdk.Coins{sdk.NewInt64Coin("ukex", 1000000)}
+	err = suite.app.BankKeeper.MintCoins(suite.ctx, minttypes.ModuleName, coins)
+	suite.Require().NoError(err)
+	err = suite.app.BankKeeper.SendCoinsFromModuleToAccount(suite.ctx, minttypes.ModuleName, addr1, coins)
+	suite.Require().NoError(err)
+	msgServer := multistakingkeeper.NewMsgServerImpl(suite.app.MultiStakingKeeper, suite.app.BankKeeper, suite.app.CustomGovKeeper, suite.app.CustomStakingKeeper)
+	_, err = msgServer.Delegate(sdk.WrapSDKContext(suite.ctx), &multistakingtypes.MsgDelegate{
+		DelegatorAddress: addr1.String(),
+		ValidatorAddress: valAddr.String(),
+		Amounts:          coins,
+	})
+	suite.Require().NoError(err)
+
+	_, err = msgServer.Undelegate(sdk.WrapSDKContext(suite.ctx), &multistakingtypes.MsgUndelegate{
+		DelegatorAddress: addr1.String(),
+		ValidatorAddress: valAddr.String(),
+		Amounts:          coins,
+	})
+	suite.Require().NoError(err)
+
+	undelegations := suite.app.MultiStakingKeeper.GetAllUndelegations(suite.ctx)
+	suite.Require().Len(undelegations, 1)
+	suite.Require().Equal(undelegations[0], multistakingtypes.Undelegation{
+		Id:         1,
+		Address:    addr1.String(),
+		ValAddress: valAddr.String(),
+		Expiry:     uint64(suite.ctx.BlockTime().Unix()) + properties.UnstakingPeriod,
+		Amount:     coins,
+	})
+
+	// claim before expiry
+	_, err = msgServer.ClaimUndelegation(sdk.WrapSDKContext(suite.ctx), &multistakingtypes.MsgClaimUndelegation{
+		Sender:         addr1.String(),
+		UndelegationId: 1,
+	})
+	suite.Require().Error(err)
+
+	// claim correctly
+	suite.ctx = suite.ctx.WithBlockTime(suite.ctx.BlockTime().Add(time.Second * time.Duration(properties.UnstakingPeriod)))
+	_, err = msgServer.ClaimUndelegation(sdk.WrapSDKContext(suite.ctx), &multistakingtypes.MsgClaimUndelegation{
+		Sender:         addr1.String(),
+		UndelegationId: 1,
+	})
+	suite.Require().NoError(err)
+
+	// Try claim again
+	_, err = msgServer.ClaimUndelegation(sdk.WrapSDKContext(suite.ctx), &multistakingtypes.MsgClaimUndelegation{
+		Sender:         addr1.String(),
+		UndelegationId: 1,
+	})
+	suite.Require().Error(err)
+}