Refactor lock method (#1936)

* Add lock method refactor * Delete duplciated testing * Update x/lockup/keeper/lock.go Co-authored-by: Aleksandr Bezobchuk <[email protected]> * Add tests implement feedback from code review * Add test cases Co-authored-by: Aleksandr Bezobchuk <[email protected]> (cherry picked from commit c115593) # Conflicts: # x/lockup/keeper/admin_keeper_test.go # x/lockup/keeper/lock.go
osmosis-labs · Jul 12, 2022 · f24bf7e · f24bf7e
1 parent 6d37b96
commit f24bf7e
Show file tree

Hide file tree

Showing 4 changed files with 181 additions and 42 deletions.
diff --git a/x/lockup/keeper/admin_keeper_test.go b/x/lockup/keeper/admin_keeper_test.go
@@ -3,8 +3,12 @@ package keeper_test
 import (
 	"time"
 
+<<<<<<< HEAD
 	"github.com/osmosis-labs/osmosis/v10/x/lockup/keeper"
 	"github.com/osmosis-labs/osmosis/v10/x/lockup/types"
+=======
+	"github.com/osmosis-labs/osmosis/v7/x/lockup/keeper"
+>>>>>>> c1155937 (Refactor `lock` method (#1936))
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -14,11 +18,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 +41,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

diff --git a/x/lockup/keeper/export_test.go b/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)
+}
diff --git a/x/lockup/keeper/lock.go b/x/lockup/keeper/lock.go
@@ -125,31 +125,41 @@ func (k Keeper) AddToExistingLock(ctx sdk.Context, owner sdk.AccAddress, coin sd
 	return locks, nil
 }
 
+<<<<<<< HEAD
 // AddTokensToLock locks more tokens into a lockup
 // This also saves the lock to the store.
 func (k Keeper) AddTokensToLockByID(ctx sdk.Context, lockID uint64, owner sdk.AccAddress, coin sdk.Coin) (*types.PeriodLock, error) {
+=======
+// AddTokensToLock locks additional tokens into an existing lock with the given ID.
+// 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, tokensToAdd sdk.Coin) (*types.PeriodLock, error) {
+>>>>>>> c1155937 (Refactor `lock` method (#1936))
 	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
 	}
 
+<<<<<<< HEAD
 	k.hooks.OnTokenLocked(ctx, lock.OwnerAddress(), lock.ID, sdk.Coins{coin}, lock.Duration, lock.EndTime)
 	return lock, nil
 }
@@ -181,22 +191,69 @@ func (k Keeper) SlashTokensFromLockByID(ctx sdk.Context, lockID uint64, coins sd
 }
 
 // LockTokens lock tokens from an account for specified duration.
+=======
+	k.hooks.AfterAddTokensToLock(ctx, lock.OwnerAddress(), lock.GetID(), sdk.NewCoins(tokensToAdd))
+
+	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
+>>>>>>> c1155937 (Refactor `lock` method (#1936))
 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 set at the beginning of unlocking time
 	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
 }
 
+<<<<<<< HEAD
 func (k Keeper) clearKeysByPrefix(ctx sdk.Context, prefix []byte) {
 	store := ctx.KVStore(k.storeKey)
 	iterator := sdk.KVStorePrefixIterator(store, prefix)
 	defer iterator.Close()
+=======
+// 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, tokensToLock); err != nil {
+		return err
+	}
+
+	// store lock object into the store
+	err = k.setLock(ctx, lock)
+	if err != nil {
+		return err
+	}
+
+	// add to accumulation store
+	for _, coin := range tokensToLock {
+		k.accumulationStore(ctx, coin.Denom).Increase(accumulationKey(lock.Duration), coin.Amount)
+	}
+
+	k.hooks.OnTokenLocked(ctx, owner, lock.ID, lock.Coins, lock.Duration, lock.EndTime)
+	return nil
+}
+>>>>>>> c1155937 (Refactor `lock` method (#1936))
 
 	for ; iterator.Valid(); iterator.Next() {
 		store.Delete(iterator.Key())

diff --git a/x/lockup/keeper/lock_test.go b/x/lockup/keeper/lock_test.go
@@ -144,47 +144,17 @@ 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()
 	now := suite.Ctx.BlockTime()
 
 	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()