Skip to content

Commit

Permalink
Refactor lock method (#1936)
Browse files Browse the repository at this point in the history 
* 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]>
  • Loading branch information
@mattverse @alexanderbez
mattverse and alexanderbez authored Jul 5, 2022
1 parent cc82980 commit c115593
Show file tree
Hide file tree
Showing 4 changed files with 141 additions and 85 deletions.
8 changes: 3 additions & 5 deletions x/lockup/keeper/admin_keeper_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
"time"

"github.com/osmosis-labs/osmosis/v7/x/lockup/keeper"
"github.com/osmosis-labs/osmosis/v7/x/lockup/types"

sdk "github.com/cosmos/cosmos-sdk/types"
)
Expand All @@ -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
Expand All @@ -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
Expand Down
4 changes: 4 additions & 0 deletions x/lockup/keeper/export_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,7 @@ func (k Keeper) GetLockRefs(ctx sdk.Context, key []byte) []uint64 {
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)
}
77 changes: 28 additions & 49 deletions x/lockup/keeper/lock.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
137 changes: 106 additions & 31 deletions x/lockup/keeper/lock_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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()

Expand Down

0 comments on commit c115593

Please sign in to comment.