Return lock id from BeginForceUnlock (#4059) (#4078)

* return lock id from BeginForceUnlock

* refactor and comment to BeginForceUnlock

* unit test BeginForceUnlock

* table driven test BeginForceUnlock

* update unit test for BeginForceUnlock

(cherry picked from commit 6ca0a56)

Co-authored-by: t4sk <[email protected]>

x/lockup/keeper/lock.go

@@ -172,33 +172,42 @@ func (k Keeper) BeginUnlock(ctx sdk.Context, lockID uint64, coins sdk.Coins) err
 		return fmt.Errorf("cannot BeginUnlocking a lock with synthetic lockup")
 	}
 
-	return k.beginUnlock(ctx, *lock, coins)
+	_, err = k.beginUnlock(ctx, *lock, coins)
+	return err
 }
 
 // BeginForceUnlock begins force unlock of the given lock.
 // This method should be called by the superfluid module ONLY, as it does not check whether
 // the lock has a synthetic lock or not before unlocking.
-func (k Keeper) BeginForceUnlock(ctx sdk.Context, lockID uint64, coins sdk.Coins) error {
+// Returns lock id, new lock id if the lock was split, else same lock id.
+func (k Keeper) BeginForceUnlock(ctx sdk.Context, lockID uint64, coins sdk.Coins) (uint64, error) {
 	lock, err := k.GetLockByID(ctx, lockID)
 	if err != nil {
-		return err
+		return 0, err
+	}
+
+	lockID, err = k.beginUnlock(ctx, *lock, coins)
+	if err != nil {
+		return 0, err
 	}
-	return k.beginUnlock(ctx, *lock, coins)
+
+	return lockID, nil
 }
 
 // beginUnlock unlocks specified tokens from the given lock. Existing lock refs
 // of not unlocking queue are deleted and new lock refs are then added.
 // EndTime of the lock is set within this method.
 // Coins provided as the parameter does not require to have all the tokens in the lock,
 // as we allow partial unlockings of a lock.
-func (k Keeper) beginUnlock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Coins) error {
+// Returns lock id, new lock id if the lock was split, else same lock id.
+func (k Keeper) beginUnlock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Coins) (uint64, error) {
 	// sanity check
 	if !coins.IsAllLTE(lock.Coins) {
-		return fmt.Errorf("requested amount to unlock exceeds locked tokens")
+		return 0, fmt.Errorf("requested amount to unlock exceeds locked tokens")
 	}
 
 	if lock.IsUnlocking() {
-		return fmt.Errorf("trying to unlock a lock that is already unlocking")
+		return 0, fmt.Errorf("trying to unlock a lock that is already unlocking")
 	}
 
 	// If the amount were unlocking is empty, or the entire coins amount, unlock the entire lock.
@@ -207,28 +216,28 @@ func (k Keeper) beginUnlock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Co
 	if len(coins) != 0 && !coins.IsEqual(lock.Coins) {
 		splitLock, err := k.splitLock(ctx, lock, coins, false)
 		if err != nil {
-			return err
+			return 0, err
 		}
 		lock = splitLock
 	}
 
 	// remove existing lock refs from not unlocking queue
 	err := k.deleteLockRefs(ctx, types.KeyPrefixNotUnlocking, lock)
 	if err != nil {
-		return err
+		return 0, err
 	}
 
 	// store lock with the end time set to current block time + duration
 	lock.EndTime = ctx.BlockTime().Add(lock.Duration)
 	err = k.setLock(ctx, lock)
 	if err != nil {
-		return err
+		return 0, err
 	}
 
 	// add lock refs into unlocking queue
 	err = k.addLockRefs(ctx, lock)
 	if err != nil {
-		return err
+		return 0, err
 	}
 
 	if k.hooks != nil {
@@ -239,7 +248,7 @@ func (k Keeper) beginUnlock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Co
 		createBeginUnlockEvent(&lock),
 	})
 
-	return nil
+	return lock.ID, nil
 }
 
 func (k Keeper) clearKeysByPrefix(ctx sdk.Context, prefix []byte) {

x/lockup/keeper/lock_test.go

@@ -45,6 +45,70 @@ func (suite *KeeperTestSuite) TestBeginUnlocking() { // test for all unlockable
 	suite.Require().NotEqual(locks[0].IsUnlocking(), false)
 }
 
+func (suite *KeeperTestSuite) TestBeginForceUnlock() {
+	// coins to lock
+	coins := sdk.NewCoins(sdk.NewInt64Coin("stake", 10))
+
+	testCases := []struct {
+		name             string
+		coins            sdk.Coins
+		unlockCoins      sdk.Coins
+		expectSameLockID bool
+	}{
+		{
+			name:             "new lock id is returned if the lock was split",
+			coins:            coins,
+			unlockCoins:      sdk.NewCoins(sdk.NewInt64Coin("stake", 1)),
+			expectSameLockID: false,
+		},
+		{
+			name:             "same lock id is returned if the lock was not split",
+			coins:            coins,
+			unlockCoins:      sdk.Coins{},
+			expectSameLockID: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(tc.name, func() {
+			suite.SetupTest()
+
+			// initial check
+			locks, err := suite.App.LockupKeeper.GetPeriodLocks(suite.Ctx)
+			suite.Require().NoError(err)
+			suite.Require().Len(locks, 0)
+
+			// lock coins
+			addr1 := sdk.AccAddress([]byte("addr1---------------"))
+			suite.LockTokens(addr1, tc.coins, time.Second)
+
+			// check locks
+			locks, err = suite.App.LockupKeeper.GetPeriodLocks(suite.Ctx)
+			suite.Require().NoError(err)
+			suite.Require().True(len(locks) > 0)
+
+			for _, lock := range locks {
+				suite.Require().Equal(lock.EndTime, time.Time{})
+				suite.Require().Equal(lock.IsUnlocking(), false)
+
+				lockID, err := suite.App.LockupKeeper.BeginForceUnlock(suite.Ctx, lock.ID, tc.unlockCoins)
+				suite.Require().NoError(err)
+
+				if tc.expectSameLockID {
+					suite.Require().Equal(lockID, lock.ID)
+				} else {
+					suite.Require().Equal(lockID, lock.ID + 1)
+				}
+
+				// new or updated lock
+				newLock, err := suite.App.LockupKeeper.GetLockByID(suite.Ctx, lockID)
+				suite.Require().NoError(err)
+				suite.Require().True(newLock.IsUnlocking())
+			}
+		})
+	}
+}
+
 func (suite *KeeperTestSuite) TestGetPeriodLocks() {
 	suite.SetupTest()
 

x/superfluid/keeper/stake.go

@@ -284,7 +284,8 @@ func (k Keeper) SuperfluidUnbondLock(ctx sdk.Context, underlyingLockId uint64, s
 	if !synthLocks[0].IsUnlocking() {
 		return types.ErrBondingLockupNotSupported
 	}
-	return k.lk.BeginForceUnlock(ctx, underlyingLockId, sdk.Coins{})
+	_, err = k.lk.BeginForceUnlock(ctx, underlyingLockId, sdk.Coins{})
+	return err
 }
 
 // alreadySuperfluidStaking returns true if underlying lock used in superfluid staking.

x/superfluid/keeper/unpool.go

@@ -77,7 +77,7 @@ func (k Keeper) UnpoolAllowedPools(ctx sdk.Context, sender sdk.AccAddress, poolI
 
 	// 8) Begin unlocking every new lock
 	for _, newLock := range newLocks {
-		err = k.lk.BeginForceUnlock(ctx, newLock.ID, newLock.Coins)
+		_, err = k.lk.BeginForceUnlock(ctx, newLock.ID, newLock.Coins)
 		if err != nil {
 			return []uint64{}, err
 		}

x/superfluid/types/expected_keepers.go

@@ -24,7 +24,7 @@ type LockupKeeper interface {
 	GetLockByID(ctx sdk.Context, lockID uint64) (*lockuptypes.PeriodLock, error)
 	// Despite the name, BeginForceUnlock is really BeginUnlock
 	// TODO: Fix this in future code update
-	BeginForceUnlock(ctx sdk.Context, lockID uint64, coins sdk.Coins) error
+	BeginForceUnlock(ctx sdk.Context, lockID uint64, coins sdk.Coins) (uint64, error)
 	ForceUnlock(ctx sdk.Context, lock lockuptypes.PeriodLock) error
 
 	CreateLock(ctx sdk.Context, owner sdk.AccAddress, coins sdk.Coins, duration time.Duration) (lockuptypes.PeriodLock, error)