Refactor `lock` method (backport #1936) (#2029)

* Refactor `lock` method (#1936)

* Add lock method refactor

* Delete duplciated testing

* Update x/lockup/keeper/lock.go

Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>

* Add tests implement feedback from code review

* Add test cases

Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>
(cherry picked from commit c1155937

)

# Conflicts:
#	x/lockup/keeper/admin_keeper_test.go
#	x/lockup/keeper/lock.go

* Fix merge conflict

* Merge

Co-authored-by: Matt, Park <45252226+mattverse@users.noreply.github.com>
Co-authored-by: mattverse <mattpark1028@gmail.com>

x/lockup/keeper/admin_keeper_test.go

@@ -4,7 +4,6 @@ import (
 	"time"
 
 	"github.com/osmosis-labs/osmosis/v10/x/lockup/keeper"
-	"github.com/osmosis-labs/osmosis/v10/x/lockup/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -14,11 +13,10 @@ func (suite *KeeperTestSuite) TestRelock() {
 
 	addr1 := sdk.AccAddress([]byte("addr1---------------"))
 	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
-	lock := types.NewPeriodLock(1, addr1, time.Second, suite.Ctx.BlockTime().Add(time.Second), coins)
 
 	// lock with balance
 	suite.FundAcc(addr1, coins)
-	err := suite.App.LockupKeeper.Lock(suite.Ctx, lock)
+	lock, err := suite.App.LockupKeeper.CreateLock(suite.Ctx, addr1, coins, time.Second)
 	suite.Require().NoError(err)
 
 	// lock with balance with same id
@@ -38,12 +36,12 @@ func (suite *KeeperTestSuite) BreakLock() {
 
 	addr1 := sdk.AccAddress([]byte("addr1---------------"))
 	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
-	lock := types.NewPeriodLock(1, addr1, time.Second, suite.Ctx.BlockTime().Add(time.Second), coins)
 
 	// lock with balance
 	suite.FundAcc(addr1, coins)
 
-	err := suite.App.LockupKeeper.Lock(suite.Ctx, lock)
+	lock, err := suite.App.LockupKeeper.CreateLock(suite.Ctx, addr1, coins, time.Second)
+
 	suite.Require().NoError(err)
 
 	// break lock

x/lockup/keeper/export_test.go

@@ -25,3 +25,7 @@ func (k Keeper) SyntheticCoins(coins sdk.Coins, suffix string) sdk.Coins {
 func (k Keeper) GetCoinsFromLocks(locks []types.PeriodLock) sdk.Coins {
 	return k.getCoinsFromLocks(locks)
 }
+
+func (k Keeper) Lock(ctx sdk.Context, lock types.PeriodLock, tokensToLock sdk.Coins) error {
+	return k.lock(ctx, lock, tokensToLock)
+}

x/lockup/keeper/lock.go

@@ -66,100 +66,79 @@ func (k Keeper) AddToExistingLock(ctx sdk.Context, owner sdk.AccAddress, coin sd
 // Tokens locked are sent and kept in the module account.
 // This method alters the lock state in store, thus we do a sanity check to ensure
 // lock owner matches the given owner.
-func (k Keeper) AddTokensToLockByID(ctx sdk.Context, lockID uint64, owner sdk.AccAddress, coin sdk.Coin) (*types.PeriodLock, error) {
+func (k Keeper) AddTokensToLockByID(ctx sdk.Context, lockID uint64, owner sdk.AccAddress, tokensToAdd sdk.Coin) (*types.PeriodLock, error) {
 	lock, err := k.GetLockByID(ctx, lockID)
+	if err != nil {
+		return nil, err
+	}
 
 	if lock.GetOwner() != owner.String() {
 		return nil, types.ErrNotLockOwner
 	}
 
+	lock.Coins = lock.Coins.Add(tokensToAdd)
+	err = k.lock(ctx, *lock, sdk.NewCoins(tokensToAdd))
 	if err != nil {
 		return nil, err
 	}
-	if err := k.bk.SendCoinsFromAccountToModule(ctx, lock.OwnerAddress(), types.ModuleName, sdk.NewCoins(coin)); err != nil {
-		return nil, err
-	}
 
-	err = k.addTokenToLock(ctx, lock, coin)
-	if err != nil {
-		return nil, err
+	for _, synthlock := range k.GetAllSyntheticLockupsByLockup(ctx, lock.ID) {
+		k.accumulationStore(ctx, synthlock.SynthDenom).Increase(accumulationKey(synthlock.Duration), tokensToAdd.Amount)
 	}
 
 	if k.hooks == nil {
 		return lock, nil
 	}
 
-	k.hooks.OnTokenLocked(ctx, lock.OwnerAddress(), lock.ID, sdk.Coins{coin}, lock.Duration, lock.EndTime)
-	return lock, nil
-}
-
-// addTokenToLock adds token to lock and modifies the state of the lock and the accumulation store.
-func (k Keeper) addTokenToLock(ctx sdk.Context, lock *types.PeriodLock, coin sdk.Coin) error {
-	lock.Coins = lock.Coins.Add(coin)
-
-	err := k.setLock(ctx, *lock)
-	if err != nil {
-		return err
-	}
-
-	// modifications to accumulation store
-	k.accumulationStore(ctx, coin.Denom).Increase(accumulationKey(lock.Duration), coin.Amount)
-
-	// CONTRACT: lock will have synthetic lock only if it has a single coin
-	// accumulation store for its synthetic denom is increased if exists.
-	lockedCoin, err := lock.SingleCoin()
-	if err == nil {
-		for _, synthlock := range k.GetAllSyntheticLockupsByLockup(ctx, lock.ID) {
-			k.accumulationStore(ctx, synthlock.SynthDenom).Increase(accumulationKey(synthlock.Duration), sdk.NewCoins(coin).AmountOf(lockedCoin.Denom))
-		}
-	}
-
-	k.hooks.AfterAddTokensToLock(ctx, lock.OwnerAddress(), lock.GetID(), sdk.NewCoins(coin))
+	k.hooks.AfterAddTokensToLock(ctx, lock.OwnerAddress(), lock.GetID(), sdk.NewCoins(tokensToAdd))
 
-	return nil
+	return lock, nil
 }
 
 // CreateLock creates a new lock with the specified duration for the owner.
+// Returns an error in the following conditions:
+//  - account does not have enough balance
 func (k Keeper) CreateLock(ctx sdk.Context, owner sdk.AccAddress, coins sdk.Coins, duration time.Duration) (types.PeriodLock, error) {
 	ID := k.GetLastLockID(ctx) + 1
 	// unlock time is initially set without a value, gets set as unlock start time + duration
 	// when unlocking starts.
 	lock := types.NewPeriodLock(ID, owner, duration, time.Time{}, coins)
-	err := k.Lock(ctx, lock)
+	err := k.lock(ctx, lock, lock.Coins)
+	if err != nil {
+		return lock, err
+	}
+
+	// add lock refs into not unlocking queue
+	err = k.addLockRefs(ctx, lock)
 	if err != nil {
 		return lock, err
 	}
+
 	k.SetLastLockID(ctx, lock.ID)
 	return lock, nil
 }
 
-// Lock is a utility method to lock tokens into the module account. This method includes setting the
-// lock within the state machine and increasing the value of accumulation store.
-func (k Keeper) Lock(ctx sdk.Context, lock types.PeriodLock) error {
+// lock is an internal utility to lock coins and set corresponding states.
+// This is only called by either of the two possible entry points to lock tokens.
+// 1. CreateLock
+// 2. AddTokensToLockByID
+func (k Keeper) lock(ctx sdk.Context, lock types.PeriodLock, tokensToLock sdk.Coins) error {
 	owner, err := sdk.AccAddressFromBech32(lock.Owner)
 	if err != nil {
 		return err
 	}
-	if err := k.bk.SendCoinsFromAccountToModule(ctx, owner, types.ModuleName, lock.Coins); err != nil {
+	if err := k.bk.SendCoinsFromAccountToModule(ctx, owner, types.ModuleName, tokensToLock); err != nil {
 		return err
 	}
 
 	// store lock object into the store
-	store := ctx.KVStore(k.storeKey)
-	bz, err := proto.Marshal(&lock)
-	if err != nil {
-		return err
-	}
-	store.Set(lockStoreKey(lock.ID), bz)
-
-	// add lock refs into not unlocking queue
-	err = k.addLockRefs(ctx, lock)
+	err = k.setLock(ctx, lock)
 	if err != nil {
 		return err
 	}
 
 	// add to accumulation store
-	for _, coin := range lock.Coins {
+	for _, coin := range tokensToLock {
 		k.accumulationStore(ctx, coin.Denom).Increase(accumulationKey(lock.Duration), coin.Amount)
 	}
 

x/lockup/keeper/lock_test.go

@@ -144,35 +144,6 @@ func (suite *KeeperTestSuite) TestUnlockPeriodLockByID() {
 	suite.Require().Len(locks, 0)
 }
 
-func (suite *KeeperTestSuite) TestLock() {
-	// test for coin locking
-	suite.SetupTest()
-
-	addr1 := sdk.AccAddress([]byte("addr1---------------"))
-	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
-	lock := types.NewPeriodLock(1, addr1, time.Second, suite.Ctx.BlockTime().Add(time.Second), coins)
-
-	// try lock without balance
-	err := suite.App.LockupKeeper.Lock(suite.Ctx, lock)
-	suite.Require().Error(err)
-
-	// lock with balance
-	suite.FundAcc(addr1, coins)
-	err = suite.App.LockupKeeper.Lock(suite.Ctx, lock)
-	suite.Require().NoError(err)
-
-	// lock with balance with same id
-	suite.FundAcc(addr1, coins)
-	err = suite.App.LockupKeeper.Lock(suite.Ctx, lock)
-	suite.Require().Error(err)
-
-	// lock with balance with different id
-	lock = types.NewPeriodLock(2, addr1, time.Second, suite.Ctx.BlockTime().Add(time.Second), coins)
-	suite.FundAcc(addr1, coins)
-	err = suite.App.LockupKeeper.Lock(suite.Ctx, lock)
-	suite.Require().NoError(err)
-}
-
 func (suite *KeeperTestSuite) TestUnlock() {
 	// test for coin unlocking
 	suite.SetupTest()
@@ -180,11 +151,10 @@ func (suite *KeeperTestSuite) TestUnlock() {
 
 	addr1 := sdk.AccAddress([]byte("addr1---------------"))
 	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
-	lock := types.NewPeriodLock(1, addr1, time.Second, time.Time{}, coins)
 
 	// lock with balance
 	suite.FundAcc(addr1, coins)
-	err := suite.App.LockupKeeper.Lock(suite.Ctx, lock)
+	lock, err := suite.App.LockupKeeper.CreateLock(suite.Ctx, addr1, coins, time.Second)
 	suite.Require().NoError(err)
 
 	// begin unlock with lock object
@@ -320,6 +290,65 @@ func (suite *KeeperTestSuite) TestLocksLongerThanDurationDenom() {
 	suite.Require().Len(locks, 1)
 }
 
+func (suite *KeeperTestSuite) TestCreateLock() {
+	suite.SetupTest()
+
+	addr1 := sdk.AccAddress([]byte("addr1---------------"))
+	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
+
+	// test locking without balance
+	_, err := suite.App.LockupKeeper.CreateLock(suite.Ctx, addr1, coins, time.Second)
+	suite.Require().Error(err)
+
+	suite.FundAcc(addr1, coins)
+
+	lock, err := suite.App.LockupKeeper.CreateLock(suite.Ctx, addr1, coins, time.Second)
+	suite.Require().NoError(err)
+
+	// check new lock
+	suite.Require().Equal(coins, lock.Coins)
+	suite.Require().Equal(time.Second, lock.Duration)
+	suite.Require().Equal(time.Time{}, lock.EndTime)
+	suite.Require().Equal(uint64(1), lock.ID)
+
+	lockID := suite.App.LockupKeeper.GetLastLockID(suite.Ctx)
+	suite.Require().Equal(uint64(1), lockID)
+
+	// check accumulation store
+	accum := suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
+		LockQueryType: types.ByDuration,
+		Denom:         "stake",
+		Duration:      time.Second,
+	})
+	suite.Require().Equal(accum.String(), "10")
+
+	// create new lock
+	coins = sdk.Coins{sdk.NewInt64Coin("stake", 20)}
+	suite.FundAcc(addr1, coins)
+
+	lock, err = suite.App.LockupKeeper.CreateLock(suite.Ctx, addr1, coins, time.Second)
+	suite.Require().NoError(err)
+
+	lockID = suite.App.LockupKeeper.GetLastLockID(suite.Ctx)
+	suite.Require().Equal(uint64(2), lockID)
+
+	// check accumulation store
+	accum = suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
+		LockQueryType: types.ByDuration,
+		Denom:         "stake",
+		Duration:      time.Second,
+	})
+	suite.Require().Equal(accum.String(), "30")
+
+	// check balance
+	balance := suite.App.BankKeeper.GetBalance(suite.Ctx, addr1, "stake")
+	suite.Require().Equal(sdk.ZeroInt(), balance.Amount)
+
+	acc := suite.App.AccountKeeper.GetModuleAccount(suite.Ctx, types.ModuleName)
+	balance = suite.App.BankKeeper.GetBalance(suite.Ctx, acc.GetAddress(), "stake")
+	suite.Require().Equal(sdk.NewInt(30), balance.Amount)
+}
+
 func (suite *KeeperTestSuite) TestAddTokensToLock() {
 	suite.SetupTest()
 
@@ -381,6 +410,52 @@ func (suite *KeeperTestSuite) TestAddTokensToLock() {
 	suite.Require().Error(err)
 }
 
+func (suite *KeeperTestSuite) TestLock() {
+	suite.SetupTest()
+
+	addr1 := sdk.AccAddress([]byte("addr1---------------"))
+	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
+
+	lock := types.PeriodLock{
+		ID:       1,
+		Owner:    addr1.String(),
+		Duration: time.Second,
+		EndTime:  time.Time{},
+		Coins:    coins,
+	}
+
+	// test locking without balance
+	err := suite.App.LockupKeeper.Lock(suite.Ctx, lock, coins)
+	suite.Require().Error(err)
+
+	// check accumulation store
+	accum := suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
+		LockQueryType: types.ByDuration,
+		Denom:         "stake",
+		Duration:      time.Second,
+	})
+	suite.Require().Equal(accum.String(), "0")
+
+	suite.FundAcc(addr1, coins)
+	err = suite.App.LockupKeeper.Lock(suite.Ctx, lock, coins)
+	suite.Require().NoError(err)
+
+	// check accumulation store
+	accum = suite.App.LockupKeeper.GetPeriodLocksAccumulation(suite.Ctx, types.QueryCondition{
+		LockQueryType: types.ByDuration,
+		Denom:         "stake",
+		Duration:      time.Second,
+	})
+	suite.Require().Equal(accum.String(), "10")
+
+	balance := suite.App.BankKeeper.GetBalance(suite.Ctx, addr1, "stake")
+	suite.Require().Equal(sdk.ZeroInt(), balance.Amount)
+
+	acc := suite.App.AccountKeeper.GetModuleAccount(suite.Ctx, types.ModuleName)
+	balance = suite.App.BankKeeper.GetBalance(suite.Ctx, acc.GetAddress(), "stake")
+	suite.Require().Equal(sdk.NewInt(10), balance.Amount)
+}
+
 func (suite *KeeperTestSuite) AddTokensToLockForSynth() {
 	suite.SetupTest()