Create Message for Staking Conversion (#5949) (#6094)

* WIP * Almost done, need more test cases... * Finish adding test cases * Use val set * Add changelog * Fix simple problems * Update tx.pb.go * Remove partial unlocking * Use val set delegation and edit test cases * Update proto/osmosis/superfluid/tx.proto Co-authored-by: Adam Tucker <[email protected]> * Update x/superfluid/keeper/stake.go Co-authored-by: alpo <[email protected]> * Romans comment, add balancer check * Adam's comment * Change errors from bool to check error * Add validate basic, bring back partial share migration for liquid gamm shares * Try abstracting check logics from test * Romans comment * Add cli * Fix cli * Update x/superfluid/keeper/stake.go * Update x/superfluid/client/cli/tx.go * Update x/superfluid/keeper/stake.go Co-authored-by: Roman <[email protected]> * Roman and Adam * Update x/superfluid/keeper/stake.go * Adams comment * Update x/superfluid/keeper/stake.go Co-authored-by: Roman <[email protected]> * Romans comments --------- Co-authored-by: devbot-wizard <[email protected]> Co-authored-by: Dev Ojha <[email protected]> Co-authored-by: Adam Tucker <[email protected]> Co-authored-by: alpo <[email protected]> Co-authored-by: Adam Tucker <[email protected]> Co-authored-by: Roman <[email protected]> (cherry picked from commit d99138f) Co-authored-by: Matt, Park <[email protected]>
osmosis-labs · Aug 17, 2023 · e4f2c3e · e4f2c3e
1 parent 88d837e
commit e4f2c3e
Show file tree

Hide file tree

Showing 20 changed files with 1,942 additions and 132 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -69,6 +69,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#5874](https://github.com/osmosis-labs/osmosis/pull/5874) Remove Partial Migration from superfluid migration to CL
 * [#5901](https://github.com/osmosis-labs/osmosis/pull/5901) Adding support for CW pools in ProtoRev
 * [#5937](https://github.com/osmosis-labs/osmosis/pull/5937) feat: add SetScalingFactorController gov prop
+* [#5949](https://github.com/osmosis-labs/osmosis/pull/5949) Add message to convert from superfluid / locks to native staking directly.
 * [#5939](https://github.com/osmosis-labs/osmosis/pull/5939) Fix: Flip existing twapRecords base/quote price denoms 
 * [#5938](https://github.com/osmosis-labs/osmosis/pull/5938) Chore: Fix valset amino codec
 

diff --git a/app/keepers/keepers.go b/app/keepers/keepers.go
@@ -394,11 +394,6 @@ func (appKeepers *AppKeepers) InitNormalKeepers(
 	appKeepers.ConcentratedLiquidityKeeper.SetIncentivesKeeper(appKeepers.IncentivesKeeper)
 	appKeepers.GAMMKeeper.SetIncentivesKeeper(appKeepers.IncentivesKeeper)
 
-	appKeepers.SuperfluidKeeper = superfluidkeeper.NewKeeper(
-		appKeepers.keys[superfluidtypes.StoreKey], appKeepers.GetSubspace(superfluidtypes.ModuleName),
-		*appKeepers.AccountKeeper, appKeepers.BankKeeper, appKeepers.StakingKeeper, appKeepers.DistrKeeper, appKeepers.EpochsKeeper, appKeepers.LockupKeeper, appKeepers.GAMMKeeper, appKeepers.IncentivesKeeper,
-		lockupkeeper.NewMsgServerImpl(appKeepers.LockupKeeper), appKeepers.ConcentratedLiquidityKeeper)
-
 	mintKeeper := mintkeeper.NewKeeper(
 		appKeepers.keys[minttypes.StoreKey],
 		appKeepers.GetSubspace(minttypes.ModuleName),
@@ -446,6 +441,11 @@ func (appKeepers *AppKeepers) InitNormalKeepers(
 
 	appKeepers.ValidatorSetPreferenceKeeper = &validatorSetPreferenceKeeper
 
+	appKeepers.SuperfluidKeeper = superfluidkeeper.NewKeeper(
+		appKeepers.keys[superfluidtypes.StoreKey], appKeepers.GetSubspace(superfluidtypes.ModuleName),
+		*appKeepers.AccountKeeper, appKeepers.BankKeeper, appKeepers.StakingKeeper, appKeepers.DistrKeeper, appKeepers.EpochsKeeper, appKeepers.LockupKeeper, appKeepers.GAMMKeeper, appKeepers.IncentivesKeeper,
+		lockupkeeper.NewMsgServerImpl(appKeepers.LockupKeeper), appKeepers.ConcentratedLiquidityKeeper, appKeepers.PoolManagerKeeper, appKeepers.ValidatorSetPreferenceKeeper)
+
 	// The last arguments can contain custom message handlers, and custom query handlers,
 	// if we want to allow any custom callbacks
 	supportedFeatures := "iterator,staking,stargate,osmosis,cosmwasm_1_1,cosmwasm_1_2"

diff --git a/proto/osmosis/superfluid/tx.proto b/proto/osmosis/superfluid/tx.proto
@@ -52,6 +52,11 @@ service Msg {
   rpc AddToConcentratedLiquiditySuperfluidPosition(
       MsgAddToConcentratedLiquiditySuperfluidPosition)
       returns (MsgAddToConcentratedLiquiditySuperfluidPositionResponse);
+
+  // UnbondConvertAndStake breaks all locks / superfluid staked assets,
+  // converts them to osmo then stakes the osmo to the designated validator.
+  rpc UnbondConvertAndStake(MsgUnbondConvertAndStake)
+      returns (MsgUnbondConvertAndStakeResponse);
 }
 
 message MsgSuperfluidDelegate {
@@ -231,4 +236,40 @@ message MsgAddToConcentratedLiquiditySuperfluidPositionResponse {
     (gogoproto.nullable) = false
   ];
   uint64 lock_id = 4 [ (gogoproto.moretags) = "yaml:\"lock_id\"" ];
+}
+
+// ===================== MsgUnbondConvertAndStake
+message MsgUnbondConvertAndStake {
+  option (amino.name) = "osmosis/unbond-convert-and-stake";
+
+  // lock ID to convert and stake.
+  // lock id with 0 should be provided if converting liquid gamm shares to stake
+  uint64 lock_id = 1 [ (gogoproto.moretags) = "yaml:\"lock_id\"" ];
+  string sender = 2 [ (gogoproto.moretags) = "yaml:\"sender\"" ];
+  // validator address to delegate to.
+  // If provided empty string, we use the validators returned from
+  // valset-preference module.
+  string val_addr = 3;
+  // min_amt_to_stake indicates the minimum amount to stake after conversion
+  string min_amt_to_stake = 4 [
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.moretags) = "yaml:\"min_amt_to_stake\"",
+    (gogoproto.nullable) = false
+  ];
+  // shares_to_convert indicates shares wanted to stake.
+  // Note that this field is only used for liquid(unlocked) gamm shares.
+  // For all other cases, this field would be disregarded.
+  cosmos.base.v1beta1.Coin shares_to_convert = 5 [
+    (gogoproto.nullable) = false,
+    (gogoproto.moretags) = "yaml:\"shares_to_convert\"",
+    (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coin"
+  ];
+}
+
+message MsgUnbondConvertAndStakeResponse {
+  string total_amt_staked = 1 [
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.moretags) = "yaml:\"total_amt_staked\"",
+    (gogoproto.nullable) = false
+  ];
 }
diff --git a/x/superfluid/client/cli/tx.go b/x/superfluid/client/cli/tx.go
@@ -34,6 +34,7 @@ func GetTxCmd() *cobra.Command {
 		// NewSuperfluidRedelegateCmd(),
 		NewCmdLockAndSuperfluidDelegate(),
 		NewCmdUnPoolWhitelistedPool(),
+		NewUnbondConvertAndStake(),
 	)
 	osmocli.AddTxCmd(cmd, NewCreateFullRangePositionAndSuperfluidDelegateCmd)
 	osmocli.AddTxCmd(cmd, NewAddToConcentratedLiquiditySuperfluidPositionCmd)
@@ -423,3 +424,56 @@ func NewUnlockAndMigrateSharesToFullRangeConcentratedPositionCmd() (*osmocli.TxC
 		Example: "unlock-and-migrate-cl 10 25000000000gamm/pool/2 1000000000uosmo,10000000uion",
 	}, &types.MsgUnlockAndMigrateSharesToFullRangeConcentratedPosition{}
 }
+
+func NewUnbondConvertAndStake() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "unbond-convert-and-stake [lock-id] [valAddr] [min-amount-to-stake](optional) [shares-to-convert](optional)",
+		Short:   "instantly unbond any locked gamm shares convert them into osmo and stake",
+		Example: "unbond-convert-and-stake 10 osmo1xxx 100000uosmo",
+		Args:    cobra.MinimumNArgs(2),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			clientCtx, err := client.GetClientTxContext(cmd)
+			if err != nil {
+				return err
+			}
+
+			txf := tx.NewFactoryCLI(clientCtx, cmd.Flags()).WithTxConfig(clientCtx.TxConfig).WithAccountRetriever(clientCtx.AccountRetriever)
+
+			sender := clientCtx.GetFromAddress()
+			lockId, err := strconv.Atoi(args[0])
+			if err != nil {
+				return err
+			}
+
+			valAddr := args[1]
+
+			var minAmtToStake sdk.Int
+			// if user provided args for min amount to stake, use it. If not, use empty coin struct
+			var sharesToConvert sdk.Coin
+			if len(args) >= 3 {
+				convertedInt, ok := sdk.NewIntFromString(args[2])
+				if !ok {
+					return fmt.Errorf("Conversion for sdk.Int failed")
+				}
+				minAmtToStake = convertedInt
+				if len(args) == 4 {
+					coins, err := sdk.ParseCoinNormalized(args[3])
+					if err != nil {
+						return err
+					}
+					sharesToConvert = coins
+				}
+			} else {
+				minAmtToStake = sdk.ZeroInt()
+				sharesToConvert = sdk.Coin{}
+			}
+
+			msg := types.NewMsgUnbondConvertAndStake(sender, uint64(lockId), valAddr, minAmtToStake, sharesToConvert)
+
+			return tx.GenerateOrBroadcastTxWithFactory(clientCtx, txf, msg)
+		},
+	}
+
+	flags.AddTxFlagsToCmd(cmd)
+	return cmd
+}
diff --git a/x/superfluid/keeper/export_test.go b/x/superfluid/keeper/export_test.go
@@ -35,12 +35,12 @@ func (k Keeper) MigrateNonSuperfluidLockBalancerToConcentrated(ctx sdk.Context,
 	return k.migrateNonSuperfluidLockBalancerToConcentrated(ctx, sender, lockId, sharesToMigrate, tokenOutMins)
 }
 
-func (k Keeper) ValidateSharesToMigrateUnlockAndExitBalancerPool(ctx sdk.Context, sender sdk.AccAddress, poolIdLeaving uint64, lock *lockuptypes.PeriodLock, sharesToMigrate sdk.Coin, tokenOutMins sdk.Coins) (exitCoins sdk.Coins, err error) {
-	return k.validateSharesToMigrateUnlockAndExitBalancerPool(ctx, sender, poolIdLeaving, lock, sharesToMigrate, tokenOutMins)
+func (k Keeper) ForceUnlockAndExitBalancerPool(ctx sdk.Context, sender sdk.AccAddress, poolIdLeaving uint64, lock *lockuptypes.PeriodLock, sharesToMigrate sdk.Coin, tokenOutMins sdk.Coins, exitCoinsLengthIsTwo bool) (exitCoins sdk.Coins, err error) {
+	return k.forceUnlockAndExitBalancerPool(ctx, sender, poolIdLeaving, lock, sharesToMigrate, tokenOutMins, exitCoinsLengthIsTwo)
 }
 
-func (k Keeper) RouteMigration(ctx sdk.Context, sender sdk.AccAddress, lockId int64, sharesToMigrate sdk.Coin) (synthLockBeforeMigration lockuptypes.SyntheticLock, migrationType MigrationType, err error) {
-	return k.routeMigration(ctx, lockId)
+func (k Keeper) GetMigrationType(ctx sdk.Context, lockId int64) (synthLockBeforeMigration lockuptypes.SyntheticLock, migrationType MigrationType, err error) {
+	return k.getMigrationType(ctx, lockId)
 }
 
 func (k Keeper) ValidateMigration(ctx sdk.Context, sender sdk.AccAddress, lockId uint64, sharesToMigrate sdk.Coin) (types.MigrationPoolIDs, *lockuptypes.PeriodLock, time.Duration, error) {
@@ -62,3 +62,25 @@ func (k Keeper) GetExistingLockRemainingDuration(ctx sdk.Context, lock *lockupty
 func (k Keeper) DistributeSuperfluidGauges(ctx sdk.Context) {
 	k.distributeSuperfluidGauges(ctx)
 }
+
+func (k Keeper) ConvertLockToStake(ctx sdk.Context, sender sdk.AccAddress, valAddr string, lockId uint64,
+	minAmtToStake sdk.Int) (totalAmtConverted sdk.Int, err error) {
+	return k.convertLockToStake(ctx, sender, valAddr, lockId, minAmtToStake)
+}
+
+func (k Keeper) ConvertGammSharesToOsmoAndStake(
+	ctx sdk.Context,
+	sender sdk.AccAddress, valAddr string,
+	poolIdLeaving uint64, exitCoins sdk.Coins, minAmtToStake sdk.Int, originalSuperfluidValAddr string,
+) (totalAmtCoverted sdk.Int, err error) {
+	return k.convertGammSharesToOsmoAndStake(ctx, sender, valAddr, poolIdLeaving, exitCoins, minAmtToStake, originalSuperfluidValAddr)
+}
+
+func (k Keeper) ConvertUnlockedToStake(ctx sdk.Context, sender sdk.AccAddress, valAddr string, sharesToStake sdk.Coin,
+	minAmtToStake sdk.Int) (totalAmtConverted sdk.Int, err error) {
+	return k.convertUnlockedToStake(ctx, sender, valAddr, sharesToStake, minAmtToStake)
+}
+
+func (k Keeper) DelegateBaseOnValsetPref(ctx sdk.Context, sender sdk.AccAddress, valAddr, originalSuperfluidValAddr string, totalAmtToStake sdk.Int) error {
+	return k.delegateBaseOnValsetPref(ctx, sender, valAddr, originalSuperfluidValAddr, totalAmtToStake)
+}
diff --git a/x/superfluid/keeper/keeper.go b/x/superfluid/keeper/keeper.go
@@ -18,23 +18,25 @@ type Keeper struct {
 	storeKey   sdk.StoreKey
 	paramSpace paramtypes.Subspace
 
-	ak  authkeeper.AccountKeeper
-	bk  types.BankKeeper
-	sk  types.StakingKeeper
-	ck  types.CommunityPoolKeeper
-	ek  types.EpochKeeper
-	lk  types.LockupKeeper
-	gk  types.GammKeeper
-	ik  types.IncentivesKeeper
-	clk types.ConcentratedKeeper
+	ak   authkeeper.AccountKeeper
+	bk   types.BankKeeper
+	sk   types.StakingKeeper
+	ck   types.CommunityPoolKeeper
+	ek   types.EpochKeeper
+	lk   types.LockupKeeper
+	gk   types.GammKeeper
+	ik   types.IncentivesKeeper
+	clk  types.ConcentratedKeeper
+	pmk  types.PoolManagerKeeper
+	vspk types.ValSetPreferenceKeeper
 
 	lms types.LockupMsgServer
 }
 
 var _ govtypes.StakingKeeper = (*Keeper)(nil)
 
 // NewKeeper returns an instance of Keeper.
-func NewKeeper(storeKey sdk.StoreKey, paramSpace paramtypes.Subspace, ak authkeeper.AccountKeeper, bk types.BankKeeper, sk types.StakingKeeper, dk types.CommunityPoolKeeper, ek types.EpochKeeper, lk types.LockupKeeper, gk types.GammKeeper, ik types.IncentivesKeeper, lms types.LockupMsgServer, clk types.ConcentratedKeeper) *Keeper {
+func NewKeeper(storeKey sdk.StoreKey, paramSpace paramtypes.Subspace, ak authkeeper.AccountKeeper, bk types.BankKeeper, sk types.StakingKeeper, dk types.CommunityPoolKeeper, ek types.EpochKeeper, lk types.LockupKeeper, gk types.GammKeeper, ik types.IncentivesKeeper, lms types.LockupMsgServer, clk types.ConcentratedKeeper, pmk types.PoolManagerKeeper, vspk types.ValSetPreferenceKeeper) *Keeper {
 	// set KeyTable if it has not already been set
 	if !paramSpace.HasKeyTable() {
 		paramSpace = paramSpace.WithKeyTable(types.ParamKeyTable())
@@ -52,6 +54,8 @@ func NewKeeper(storeKey sdk.StoreKey, paramSpace paramtypes.Subspace, ak authkee
 		gk:         gk,
 		ik:         ik,
 		clk:        clk,
+		pmk:        pmk,
+		vspk:       vspk,
 
 		lms: lms,
 	}

diff --git a/x/superfluid/keeper/migrate.go b/x/superfluid/keeper/migrate.go
@@ -47,7 +47,7 @@ const (
 //
 // Errors if the lock is not found, if the lock is not a balancer pool lock, or if the lock is not owned by the sender.
 func (k Keeper) RouteLockedBalancerToConcentratedMigration(ctx sdk.Context, sender sdk.AccAddress, providedLockId int64, sharesToMigrate sdk.Coin, tokenOutMins sdk.Coins) (positionData cltypes.CreateFullRangePositionData, migratedPoolIDs types.MigrationPoolIDs, concentratedLockId uint64, err error) {
-	synthLockBeforeMigration, migrationType, err := k.routeMigration(ctx, providedLockId)
+	synthLockBeforeMigration, migrationType, err := k.getMigrationType(ctx, providedLockId)
 	if err != nil {
 		return cltypes.CreateFullRangePositionData{}, types.MigrationPoolIDs{}, 0, err
 	}
@@ -105,7 +105,7 @@ func (k Keeper) migrateSuperfluidBondedBalancerToConcentrated(ctx sdk.Context,
 	// Force unlock, validate the provided sharesToMigrate, and exit the balancer pool.
 	// This will return the coins that will be used to create the concentrated liquidity position.
 	// It also returns the lock object that contains the remaining shares that were not used in this migration.
-	exitCoins, err := k.validateSharesToMigrateUnlockAndExitBalancerPool(ctx, sender, migratedPoolIDs.LeavingID, preMigrationLock, sharesToMigrate, tokenOutMins)
+	exitCoins, err := k.forceUnlockAndExitBalancerPool(ctx, sender, migratedPoolIDs.LeavingID, preMigrationLock, sharesToMigrate, tokenOutMins, true)
 	if err != nil {
 		return cltypes.CreateFullRangePositionData{}, 0, types.MigrationPoolIDs{}, err
 	}
@@ -151,7 +151,7 @@ func (k Keeper) migrateSuperfluidUnbondingBalancerToConcentrated(ctx sdk.Context
 	// Force unlock, validate the provided sharesToMigrate, and exit the balancer pool.
 	// This will return the coins that will be used to create the concentrated liquidity position.
 	// It also returns the lock object that contains the remaining shares that were not used in this migration.
-	exitCoins, err := k.validateSharesToMigrateUnlockAndExitBalancerPool(ctx, sender, migratedPoolIDs.LeavingID, preMigrationLock, sharesToMigrate, tokenOutMins)
+	exitCoins, err := k.forceUnlockAndExitBalancerPool(ctx, sender, migratedPoolIDs.LeavingID, preMigrationLock, sharesToMigrate, tokenOutMins, true)
 	if err != nil {
 		return cltypes.CreateFullRangePositionData{}, 0, types.MigrationPoolIDs{}, err
 	}
@@ -199,7 +199,7 @@ func (k Keeper) migrateNonSuperfluidLockBalancerToConcentrated(ctx sdk.Context,
 	// Force unlock, validate the provided sharesToMigrate, and exit the balancer pool.
 	// This will return the coins that will be used to create the concentrated liquidity position.
 	// It also returns the lock object that contains the remaining shares that were not used in this migration.
-	exitCoins, err := k.validateSharesToMigrateUnlockAndExitBalancerPool(ctx, sender, migratedPoolIDs.LeavingID, preMigrationLock, sharesToMigrate, tokenOutMins)
+	exitCoins, err := k.forceUnlockAndExitBalancerPool(ctx, sender, migratedPoolIDs.LeavingID, preMigrationLock, sharesToMigrate, tokenOutMins, true)
 	if err != nil {
 		return cltypes.CreateFullRangePositionData{}, 0, types.MigrationPoolIDs{}, err
 	}
@@ -215,9 +215,9 @@ func (k Keeper) migrateNonSuperfluidLockBalancerToConcentrated(ctx sdk.Context,
 	return positionData, concentratedLockId, migratedPoolIDs, nil
 }
 
-// routeMigration determines the status of the provided lock which is used to determine the method for migration.
+// getMigrationType determines the status of the provided lock which is used to determine the method for migration.
 // It also returns the underlying synthetic locks of the provided lock, if any exist.
-func (k Keeper) routeMigration(ctx sdk.Context, providedLockId int64) (synthLockBeforeMigration lockuptypes.SyntheticLock, migrationType MigrationType, err error) {
+func (k Keeper) getMigrationType(ctx sdk.Context, providedLockId int64) (synthLockBeforeMigration lockuptypes.SyntheticLock, migrationType MigrationType, err error) {
 	// As a hack around to get frontend working, we decided to allow negative values for the provided lock ID to indicate that the user wants to migrate shares that are not locked.
 	if providedLockId <= 0 {
 		return lockuptypes.SyntheticLock{}, Unlocked, nil
@@ -294,15 +294,15 @@ func (k Keeper) validateMigration(ctx sdk.Context, sender sdk.AccAddress, lockId
 	}, preMigrationLock, remainingLockTime, nil
 }
 
-// validateSharesToMigrateUnlockAndExitBalancerPool validates the unlocking and exiting of gamm LP tokens from the Balancer pool. It performs the following steps:
+// forceUnlockAndExitBalancerPool validates the unlocking and exiting of gamm LP tokens from the Balancer pool. It performs the following steps:
 //
 // 1. Completes the unlocking process / deletes synthetic locks for the provided lock.
 // 2. If shares to migrate are not specified, all shares in the lock are migrated.
 // 3. Ensures that the number of shares to migrate is less than or equal to the number of shares in the lock.
 // 4. Exits the position in the Balancer pool.
 // 5. Ensures that exactly two coins are returned.
 // 6. Any remaining shares that were not migrated are re-locked as a new lock for the remaining time on the lock.
-func (k Keeper) validateSharesToMigrateUnlockAndExitBalancerPool(ctx sdk.Context, sender sdk.AccAddress, poolIdLeaving uint64, lock *lockuptypes.PeriodLock, sharesToMigrate sdk.Coin, tokenOutMins sdk.Coins) (exitCoins sdk.Coins, err error) {
+func (k Keeper) forceUnlockAndExitBalancerPool(ctx sdk.Context, sender sdk.AccAddress, poolIdLeaving uint64, lock *lockuptypes.PeriodLock, sharesToMigrate sdk.Coin, tokenOutMins sdk.Coins, exitCoinsLengthIsTwo bool) (exitCoins sdk.Coins, err error) {
 	// validateMigration ensures that the preMigrationLock contains coins of length 1.
 	gammSharesInLock := lock.Coins[0]
 
@@ -321,7 +321,7 @@ func (k Keeper) validateSharesToMigrateUnlockAndExitBalancerPool(ctx sdk.Context
 		return sdk.Coins{}, types.MigratePartialSharesError{SharesToMigrate: sharesToMigrate.Amount.String(), SharesInLock: gammSharesInLock.Amount.String()}
 	}
 
-	// Force migrate, which breaks and deletes associated synthetic locks.
+	// Force migrate, which breaks and deletes associated synthetic locks (if exists).
 	err = k.lk.ForceUnlock(ctx, *lock)
 	if err != nil {
 		return sdk.Coins{}, err
@@ -333,8 +333,8 @@ func (k Keeper) validateSharesToMigrateUnlockAndExitBalancerPool(ctx sdk.Context
 		return sdk.Coins{}, err
 	}
 
-	// Defense in depth, ensuring we are returning exactly two coins.
-	if len(exitCoins) != 2 {
+	// if exit coins length should be two, check exitCoins length
+	if exitCoinsLengthIsTwo && len(exitCoins) != 2 {
 		return sdk.Coins{}, types.TwoTokenBalancerPoolError{NumberOfTokens: len(exitCoins)}
 	}