chore: Refactor and Add tests for AddToExistingLock

x/lockup/keeper/lock.go

@@ -48,18 +48,22 @@ func (k Keeper) BeginUnlockAllNotUnlockings(ctx sdk.Context, account sdk.AccAddr
 }
 
 // AddToExistingLock adds the given coin to the existing lock with the same owner and duration.
-// Returns an empty array of period lock when a lock with the given condition does not exist.
-func (k Keeper) AddToExistingLock(ctx sdk.Context, owner sdk.AccAddress, coin sdk.Coin, duration time.Duration) ([]types.PeriodLock, error) {
+// Returns the updated lock ID if successfully added coin, returns 0 when a lock with
+// given condition does not exist.
+func (k Keeper) AddToExistingLock(ctx sdk.Context, owner sdk.AccAddress, coin sdk.Coin, duration time.Duration) (uint64, error) {
 	locks := k.GetAccountLockedDurationNotUnlockingOnly(ctx, owner, coin.Denom, duration)
 	// if existing lock with same duration and denom exists, just add there
 	if len(locks) > 0 {
+		// there should only be a single lock with the same duration + token, thus we take the first lock
 		lock := locks[0]
 		_, err := k.AddTokensToLockByID(ctx, lock.ID, owner, coin)
 		if err != nil {
-			return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, err.Error())
+			return 0, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, err.Error())
 		}
+		return lock.ID, nil
 	}
-	return locks, nil
+
+	return 0, nil
 }
 
 // AddTokensToLock locks additional tokens into an existing lock with the given ID.

x/lockup/keeper/lock_test.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/osmosis-labs/osmosis/v7/x/lockup/types"
 
-	"github.com/cosmos/cosmos-sdk/simapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -350,64 +349,108 @@ func (suite *KeeperTestSuite) TestCreateLock() {
 }
 
 func (suite *KeeperTestSuite) TestAddTokensToLock() {
-	suite.SetupTest()
-
-	// lock coins
+	initialLockCoin := sdk.NewInt64Coin("stake", 10)
 	addr1 := sdk.AccAddress([]byte("addr1---------------"))
-	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
-	suite.LockTokens(addr1, coins, time.Second)
-
-	// check locks
-	locks, err := suite.App.LockupKeeper.GetPeriodLocks(suite.Ctx)
-	suite.Require().NoError(err)
-	suite.Require().Len(locks, 1)
-	suite.Require().Equal(locks[0].Coins, coins)
-	// check accumulation store is correctly updated
-	accum := suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
-		LockQueryType: types.ByDuration,
-		Denom:         "stake",
-		Duration:      time.Second,
-	})
-	suite.Require().Equal(accum.String(), "10")
-
-	// add more tokens to lock
-	addCoins := sdk.NewInt64Coin("stake", 10)
-	suite.FundAcc(addr1, sdk.Coins{addCoins})
-	_, err = suite.App.LockupKeeper.AddTokensToLockByID(suite.Ctx, locks[0].ID, addr1, addCoins)
-	suite.Require().NoError(err)
+	addr2 := sdk.AccAddress([]byte("addr2---------------"))
 
-	// check locks after adding tokens to lock
-	locks, err = suite.App.LockupKeeper.GetPeriodLocks(suite.Ctx)
-	suite.Require().NoError(err)
-	suite.Require().Len(locks, 1)
-	suite.Require().Equal(locks[0].Coins, coins.Add(sdk.Coins{addCoins}...))
+	testCases := []struct {
+		name                         string
+		tokenToAdd                   sdk.Coin
+		duration                     time.Duration
+		lockingAddress               sdk.AccAddress
+		expectAddTokensToLockSuccess bool
+		expectedError                bool
+	}{
+		{
+			name:                         "normal add tokens to lock case",
+			tokenToAdd:                   initialLockCoin,
+			duration:                     time.Second,
+			lockingAddress:               addr1,
+			expectAddTokensToLockSuccess: true,
+		},
+		{
+			name:           "not the owner of the lock",
+			tokenToAdd:     initialLockCoin,
+			duration:       time.Second,
+			lockingAddress: addr2,
+		},
+		{
+			name:           "lock with matching duration not existing",
+			tokenToAdd:     initialLockCoin,
+			duration:       time.Second * 2,
+			lockingAddress: addr1,
+		},
+		{
+			name:           "lock invalid tokens",
+			tokenToAdd:     sdk.NewCoin("unknown", sdk.NewInt(10)),
+			duration:       time.Second,
+			lockingAddress: addr1,
+		},
+		{
+			name:           "token to add exceeds balance",
+			tokenToAdd:     sdk.NewCoin("stake", sdk.NewInt(20)),
+			duration:       time.Second,
+			lockingAddress: addr1,
+			expectedError:  true,
+		},
+	}
 
-	// check accumulation store is correctly updated
-	accum = suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
-		LockQueryType: types.ByDuration,
-		Denom:         "stake",
-		Duration:      time.Second,
-	})
-	suite.Require().Equal(accum.String(), "20")
+	for _, tc := range testCases {
+		suite.SetupTest()
+		// lock with balance
+		suite.FundAcc(addr1, sdk.Coins{initialLockCoin})
+		originalLock, err := suite.App.LockupKeeper.CreateLock(suite.Ctx, addr1, sdk.Coins{initialLockCoin}, time.Second)
+		suite.Require().NoError(err)
 
-	// try to add tokens to unavailable lock
-	cacheCtx, _ := suite.Ctx.CacheContext()
-	err = simapp.FundAccount(suite.App.BankKeeper, cacheCtx, addr1, sdk.Coins{addCoins})
-	suite.Require().NoError(err)
-	// curBalance := suite.App.BankKeeper.GetAllBalances(cacheCtx, addr1)
-	_, err = suite.App.LockupKeeper.AddTokensToLockByID(cacheCtx, 1111, addr1, addCoins)
-	suite.Require().Error(err)
+		suite.FundAcc(addr1, sdk.Coins{initialLockCoin})
+		balanceBeforeLock := suite.App.BankKeeper.GetAllBalances(suite.Ctx, tc.lockingAddress)
 
-	// try to add tokens with lack balance
-	cacheCtx, _ = suite.Ctx.CacheContext()
-	_, err = suite.App.LockupKeeper.AddTokensToLockByID(cacheCtx, locks[0].ID, addr1, addCoins)
-	suite.Require().Error(err)
+		lockID, err := suite.App.LockupKeeper.AddToExistingLock(suite.Ctx, tc.lockingAddress, tc.tokenToAdd, tc.duration)
+		if tc.expectedError {
+			suite.Require().Error(err)
+		} else {
+			suite.Require().NoError(err)
+		}
 
-	// try to add tokens to lock that is owned by others
-	addr2 := sdk.AccAddress([]byte("addr2---------------"))
-	suite.FundAcc(addr2, sdk.Coins{addCoins})
-	_, err = suite.App.LockupKeeper.AddTokensToLockByID(cacheCtx, locks[0].ID, addr2, addCoins)
-	suite.Require().Error(err)
+		if tc.expectAddTokensToLockSuccess {
+			suite.Require().NoError(err)
+
+			// get updated lock
+			lock, err := suite.App.LockupKeeper.GetLockByID(suite.Ctx, lockID)
+			suite.Require().NoError(err)
+
+			// check that tokens have been added successfully to the lock
+			suite.Require().True(sdk.Coins{initialLockCoin.Add(tc.tokenToAdd)}.IsEqual(lock.Coins))
+
+			// check balance has decreased
+			balanceAfterLock := suite.App.BankKeeper.GetAllBalances(suite.Ctx, tc.lockingAddress)
+			suite.Require().True(balanceBeforeLock.IsEqual(balanceAfterLock.Add(tc.tokenToAdd)))
+
+			// check accumulation store
+			accum := suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
+				LockQueryType: types.ByDuration,
+				Denom:         "stake",
+				Duration:      time.Second,
+			})
+			suite.Require().Equal(initialLockCoin.Amount.Add(tc.tokenToAdd.Amount), accum)
+		} else {
+			suite.Require().Equal(uint64(0), lockID)
+
+			lock, err := suite.App.LockupKeeper.GetLockByID(suite.Ctx, originalLock.ID)
+			suite.Require().NoError(err)
+
+			// check that locked coins haven't changed
+			suite.Require().True(lock.Coins.IsEqual(sdk.Coins{initialLockCoin}))
+
+			// check accumulation store didn't change
+			accum := suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
+				LockQueryType: types.ByDuration,
+				Denom:         "stake",
+				Duration:      time.Second,
+			})
+			suite.Require().Equal(initialLockCoin.Amount, accum)
+		}
+	}
 }
 
 func (suite *KeeperTestSuite) TestLock() {

x/lockup/keeper/msg_server.go

@@ -39,22 +39,22 @@ func (server msgServer) LockTokens(goCtx context.Context, msg *types.MsgLockToke
 
 	// check if there's an existing lock from the same owner with the same duration.
 	// If so, simply add tokens to the existing lock.
-	locks, err := server.keeper.AddToExistingLock(ctx, owner, msg.Coins[0], msg.Duration)
+	lockID, err := server.keeper.AddToExistingLock(ctx, owner, msg.Coins[0], msg.Duration)
 	if err != nil {
 		return nil, err
 	}
 
 	// return the lock id of the existing lock when successfully added to the existing lock.
-	if len(locks) > 0 {
+	if lockID != 0 {
 		ctx.EventManager().EmitEvents(sdk.Events{
 			sdk.NewEvent(
 				types.TypeEvtAddTokensToLock,
-				sdk.NewAttribute(types.AttributePeriodLockID, utils.Uint64ToString(locks[0].ID)),
+				sdk.NewAttribute(types.AttributePeriodLockID, utils.Uint64ToString(lockID)),
 				sdk.NewAttribute(types.AttributePeriodLockOwner, msg.Owner),
 				sdk.NewAttribute(types.AttributePeriodLockAmount, msg.Coins.String()),
 			),
 		})
-		return &types.MsgLockTokensResponse{ID: locks[0].ID}, nil
+		return &types.MsgLockTokensResponse{ID: lockID}, nil
 	}
 
 	// if the owner + duration combination is new, create a new lock.