added undelegate logic

x/valset-pref/msg_server.go

@@ -47,7 +47,15 @@ func (server msgServer) DelegateToValidatorSet(goCtx context.Context, msg *types
 
 	return &types.MsgDelegateToValidatorSetResponse{}, nil
 }
+
 func (server msgServer) UndelegateFromValidatorSet(goCtx context.Context, msg *types.MsgUndelegateFromValidatorSet) (*types.MsgUndelegateFromValidatorSetResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	err := server.keeper.UndelegateFromValidatorSet(ctx, msg.Delegator, msg.Coin)
+	if err != nil {
+		return nil, err
+	}
+
 	return &types.MsgUndelegateFromValidatorSetResponse{}, nil
 }
 

x/valset-pref/msg_server_test.go

@@ -202,3 +202,64 @@ func (suite *KeeperTestSuite) TestDelegateToValidatorSet() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestUnDelegateFromValidatorSet() {
+	tests := []struct {
+		name          string
+		delegator     sdk.AccAddress
+		coinToStake   sdk.Coin
+		coinToUnStake sdk.Coin
+		expectPass    bool
+	}{
+		{
+			name:          "Unstake x amount from a validator Set",
+			delegator:     sdk.AccAddress([]byte("addr1---------------")),
+			coinToStake:   sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),
+			coinToUnStake: sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(10_000_000)),
+			expectPass:    true,
+		},
+		{
+			name:          "Unstake everything",
+			delegator:     sdk.AccAddress([]byte("addr2---------------")),
+			coinToStake:   sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),
+			coinToUnStake: sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),
+			expectPass:    true,
+		},
+		{
+			name:          "Unstake more amount than the staked amount",
+			delegator:     sdk.AccAddress([]byte("addr3---------------")),
+			coinToStake:   sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),
+			coinToUnStake: sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(40_000_000)),
+			expectPass:    false,
+		},
+	}
+
+	for _, test := range tests {
+		suite.Run(test.name, func() {
+			suite.SetupTest()
+
+			suite.FundAcc(test.delegator, sdk.Coins{sdk.NewInt64Coin(sdk.DefaultBondDenom, 100_000_000)}) // 100 osmo
+
+			// setup message server
+			msgServer := valPref.NewMsgServerImpl(suite.App.ValidatorSetPreferenceKeeper)
+			c := sdk.WrapSDKContext(suite.Ctx)
+
+			// call the create validator set preference
+			preferences := suite.PrepareDelegateToValidatorSet()
+
+			_, err := msgServer.SetValidatorSetPreference(c, types.NewMsgSetValidatorSetPreference(test.delegator, preferences))
+			suite.Require().NoError(err)
+
+			// call the create validator set preference
+			_, err = msgServer.DelegateToValidatorSet(c, types.NewMsgDelegateToValidatorSet(test.delegator, test.coinToStake))
+			suite.Require().NoError(err)
+
+			_, err = msgServer.UndelegateFromValidatorSet(c, types.NewMsgUndelegateFromValidatorSet(test.delegator, test.coinToUnStake))
+			if test.expectPass {
+				suite.Require().NoError(err)
+			} else {
+				suite.Require().Error(err)
+			}
+		})
+	}
+}

x/valset-pref/validator_set.go

@@ -65,6 +65,56 @@ func (k Keeper) DelegateToValidatorSet(ctx sdk.Context, delegatorAddr string, co
 	return nil
 }
 
+// UndelegateFromValidatorSet undelegates {coin} amount from the validator set.
+// For ex: userA has staked 10tokens with weight {Val->0.5, ValB->0.3, ValC->0.2}
+// undelegate 6osmo with validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2}
+// our undelegate logic would attempt to undelegate 3osmo from A , 1.8osmo from B, 1.2osmo from C
+func (k Keeper) UndelegateFromValidatorSet(ctx sdk.Context, delegatorAddr string, coin sdk.Coin) error {
+	// get the existing validator set preference
+	existingSet, found := k.GetValidatorSetPreference(ctx, delegatorAddr)
+	if !found {
+		return fmt.Errorf("user %s doesn't have validator set", delegatorAddr)
+	}
+
+	delegator, err := sdk.AccAddressFromBech32(delegatorAddr)
+	if err != nil {
+		return err
+	}
+
+	// the total amount the user wants to undelegate
+	tokenAmt := sdk.NewDec(coin.Amount.Int64())
+
+	totalAmountFromWeights := sdk.NewDec(0)
+	for _, val := range existingSet.Preferences {
+		totalAmountFromWeights = totalAmountFromWeights.Add(val.Weight.Mul(tokenAmt))
+	}
+
+	if !totalAmountFromWeights.Equal(tokenAmt) {
+		return fmt.Errorf("The undelegate total donot add up with the amount calculated from weights expected %s got %s", tokenAmt, totalAmountFromWeights)
+	}
+
+	for _, val := range existingSet.Preferences {
+		// Calculate the amount to undelegate based on the existing weights
+		amountToUnDelegate := val.Weight.Mul(tokenAmt)
+
+		valAddr, validator, err := k.getValAddrAndVal(ctx, val.ValOperAddress)
+		if err != nil {
+			return err
+		}
+
+		sharesAmt, err := validator.SharesFromTokens(amountToUnDelegate.TruncateInt())
+		if err != nil {
+			return err
+		}
+
+		_, err = k.stakingKeeper.Undelegate(ctx, delegator, valAddr, sharesAmt) // this has to be shares amount
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // GetValAddrAndVal checks if the validator address is valid and the validator provided exists on chain.
 func (k Keeper) getValAddrAndVal(ctx sdk.Context, valOperAddress string) (sdk.ValAddress, stakingtypes.Validator, error) {
 	valAddr, err := sdk.ValAddressFromBech32(valOperAddress)