chore: Refactor and Add tests for AddToExistingLock (#1979)

* Refactor and add test for addtoexistinglock * Add documentation * Return lockid from AddToExistingLock * Add lock testing * Add HasLock * fix lint * Add error to AddToExistingLock, add changelog entry * Add test cases * Update x/lockup/keeper/lock.go Co-authored-by: Roman <[email protected]> * Remove unnecessary fields Co-authored-by: Roman <[email protected]>
osmosis-labs · Aug 10, 2022 · 08669da · 08669da
1 parent b5c98de
commit 08669da
Show file tree

Hide file tree

Showing 4 changed files with 190 additions and 68 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -50,6 +50,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   * MsgExitPoolResponse: token_out field 
 * [#1825](https://github.com/osmosis-labs/osmosis/pull/1825) Fixes Interchain Accounts (host side) by adding it to AppModuleBasics
 * [#1699](https://github.com/osmosis-labs/osmosis/pull/1699) Fixes bug in sig fig rounding on spot price queries for small values
+* [#1699](https://github.com/osmosis-labs/osmosis/pull/1979) `AddToExistingLock` returns error when lock with matching conditions does not exist.
 * [#1994](https://github.com/osmosis-labs/osmosis/pull/1994) Removed bech32ibc module
 * [#2016](https://github.com/osmosis-labs/osmosis/pull/2016) Add fixed 10000 gas cost for each Balancer swap
 * [#2147](https://github.com/osmosis-labs/osmosis/pull/2147) Set MaxAgeNumBlocks in v11 Upgrade Handler to two weeks.

diff --git a/x/lockup/keeper/lock.go b/x/lockup/keeper/lock.go
@@ -48,18 +48,31 @@ 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 and error when a lock with
+// given condition does not exist, or if fails to add to lock.
+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 {
-		lock := locks[0]
-		_, err := k.AddTokensToLockByID(ctx, lock.ID, owner, coin)
-		if err != nil {
-			return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, err.Error())
-		}
+
+	// if no lock exists for the given owner + denom + duration, return an error
+	if len(locks) < 1 {
+		return 0, sdkerrors.Wrapf(types.ErrLockupNotFound, "lock with denom %s before duration %s does not exist", coin.Denom, duration.String())
 	}
-	return locks, nil
+
+	// if existing lock with same duration and denom exists, add to the existing lock
+	// 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 0, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, err.Error())
+	}
+
+	return lock.ID, nil
+}
+
+// HasLock returns true if lock with the given condition exists
+func (k Keeper) HasLock(ctx sdk.Context, owner sdk.AccAddress, denom string, duration time.Duration) bool {
+	locks := k.GetAccountLockedDurationNotUnlockingOnly(ctx, owner, denom, duration)
+	return len(locks) > 0
 }
 
 // AddTokensToLock locks additional tokens into an existing lock with the given ID.

diff --git a/x/lockup/keeper/lock_test.go b/x/lockup/keeper/lock_test.go
@@ -6,7 +6,6 @@ import (
 
 	"github.com/osmosis-labs/osmosis/v10/x/lockup/types"
 
-	"github.com/cosmos/cosmos-sdk/simapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -333,64 +332,173 @@ 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)
+	addr2 := sdk.AccAddress([]byte("addr2---------------"))
 
-	// 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")
+	testCases := []struct {
+		name                         string
+		tokenToAdd                   sdk.Coin
+		duration                     time.Duration
+		lockingAddress               sdk.AccAddress
+		expectAddTokensToLockSuccess 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,
+		},
+	}
 
-	// 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)
+	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)
 
-	// 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}...))
+		suite.FundAcc(addr1, sdk.Coins{initialLockCoin})
+		balanceBeforeLock := suite.App.BankKeeper.GetAllBalances(suite.Ctx, tc.lockingAddress)
 
-	// 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")
+		lockID, err := suite.App.LockupKeeper.AddToExistingLock(suite.Ctx, tc.lockingAddress, tc.tokenToAdd, tc.duration)
 
-	// 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)
+		if tc.expectAddTokensToLockSuccess {
+			suite.Require().NoError(err)
 
-	// 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)
+			// 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().Error(err)
+			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)
+		}
+	}
+}
 
-	// try to add tokens to lock that is owned by others
+func (suite *KeeperTestSuite) TestHasLock() {
+	addr1 := sdk.AccAddress([]byte("addr1---------------"))
 	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)
+
+	testCases := []struct {
+		name            string
+		tokenLocked     sdk.Coin
+		durationLocked  time.Duration
+		lockAddr        sdk.AccAddress
+		denomToQuery    string
+		durationToQuery time.Duration
+		expectedHas     bool
+	}{
+		{
+			name:            "same token, same duration",
+			tokenLocked:     sdk.NewInt64Coin("stake", 10),
+			durationLocked:  time.Minute,
+			lockAddr:        addr1,
+			denomToQuery:    "stake",
+			durationToQuery: time.Minute,
+			expectedHas:     true,
+		},
+		{
+			name:            "same token, shorter duration",
+			tokenLocked:     sdk.NewInt64Coin("stake", 10),
+			durationLocked:  time.Minute,
+			lockAddr:        addr1,
+			denomToQuery:    "stake",
+			durationToQuery: time.Second,
+			expectedHas:     false,
+		},
+		{
+			name:            "same token, longer duration",
+			tokenLocked:     sdk.NewInt64Coin("stake", 10),
+			durationLocked:  time.Minute,
+			lockAddr:        addr1,
+			denomToQuery:    "stake",
+			durationToQuery: time.Minute * 2,
+			expectedHas:     false,
+		},
+		{
+			name:            "different token, same duration",
+			tokenLocked:     sdk.NewInt64Coin("stake", 10),
+			durationLocked:  time.Minute,
+			lockAddr:        addr1,
+			denomToQuery:    "uosmo",
+			durationToQuery: time.Minute,
+			expectedHas:     false,
+		},
+		{
+			name:            "same token, same duration, different address",
+			tokenLocked:     sdk.NewInt64Coin("stake", 10),
+			durationLocked:  time.Minute,
+			lockAddr:        addr2,
+			denomToQuery:    "uosmo",
+			durationToQuery: time.Minute,
+			expectedHas:     false,
+		},
+	}
+	for _, tc := range testCases {
+		suite.SetupTest()
+
+		suite.FundAcc(tc.lockAddr, sdk.Coins{tc.tokenLocked})
+		_, err := suite.App.LockupKeeper.CreateLock(suite.Ctx, tc.lockAddr, sdk.Coins{tc.tokenLocked}, tc.durationLocked)
+		suite.Require().NoError(err)
+
+		hasLock := suite.App.LockupKeeper.HasLock(suite.Ctx, addr1, tc.denomToQuery, tc.durationToQuery)
+		suite.Require().Equal(tc.expectedHas, hasLock)
+	}
 }
 
 func (suite *KeeperTestSuite) TestLock() {

diff --git a/x/lockup/keeper/msg_server.go b/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)
-	if err != nil {
-		return nil, err
-	}
+	lockExists := server.keeper.HasLock(ctx, owner, msg.Coins[0].Denom, msg.Duration)
+	if lockExists {
+		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 {
 		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.