Skip to content

Commit

Permalink
refactor: utilize collections for state management (#23)
Browse files Browse the repository at this point in the history
  • Loading branch information
johnletey authored Oct 8, 2024
1 parent ba86cbf commit 099f728
Show file tree
Hide file tree
Showing 17 changed files with 796 additions and 276 deletions.
36 changes: 27 additions & 9 deletions genesis.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,24 +9,42 @@ import (
)

func InitGenesis(ctx sdk.Context, k *keeper.Keeper, addressCodec address.Codec, genesis types.GenesisState) {
k.SetBlocklistOwner(ctx, genesis.BlocklistState.Owner)
k.SetBlocklistPendingOwner(ctx, genesis.BlocklistState.PendingOwner)
if err := k.SetBlocklistOwner(ctx, genesis.BlocklistState.Owner); err != nil {
panic(err)
}
if err := k.SetBlocklistPendingOwner(ctx, genesis.BlocklistState.PendingOwner); err != nil {
panic(err)
}
for _, account := range genesis.BlocklistState.BlockedAddresses {
address, _ := addressCodec.StringToBytes(account)
k.SetBlockedAddress(ctx, address)
if err := k.SetBlockedAddress(ctx, address); err != nil {
panic(err)
}
}

k.SetPaused(ctx, genesis.Paused)
k.SetOwner(ctx, genesis.Owner)
k.SetPendingOwner(ctx, genesis.PendingOwner)
if err := k.SetPaused(ctx, genesis.Paused); err != nil {
panic(err)
}
if err := k.SetOwner(ctx, genesis.Owner); err != nil {
panic(err)
}
if err := k.SetPendingOwner(ctx, genesis.PendingOwner); err != nil {
panic(err)
}
for _, burner := range genesis.Burners {
k.SetBurner(ctx, burner.Address, burner.Allowance)
if err := k.SetBurner(ctx, burner.Address, burner.Allowance); err != nil {
panic(err)
}
}
for _, minter := range genesis.Minters {
k.SetMinter(ctx, minter.Address, minter.Allowance)
if err := k.SetMinter(ctx, minter.Address, minter.Allowance); err != nil {
panic(err)
}
}
for _, pauser := range genesis.Pausers {
k.SetPauser(ctx, pauser)
if err := k.SetPauser(ctx, pauser); err != nil {
panic(err)
}
}
}

Expand Down
4 changes: 2 additions & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,14 @@ go 1.22.7

require (
cosmossdk.io/api v0.7.5
cosmossdk.io/collections v0.4.0
cosmossdk.io/core v0.11.1
cosmossdk.io/depinject v1.0.0
cosmossdk.io/errors v1.0.1
cosmossdk.io/math v1.3.0
cosmossdk.io/store v1.1.1
github.com/cometbft/cometbft v0.38.12
github.com/cosmos/cosmos-db v1.0.2
github.com/cosmos/cosmos-proto v1.0.0-beta.5
github.com/cosmos/cosmos-sdk v0.50.10
github.com/cosmos/gogoproto v1.7.0
Expand All @@ -27,7 +29,6 @@ require (
require (
4d63.com/gocheckcompilerdirectives v1.2.1 // indirect
4d63.com/gochecknoglobals v0.2.1 // indirect
cosmossdk.io/collections v0.4.0 // indirect
cosmossdk.io/log v1.4.1 // indirect
cosmossdk.io/x/tx v0.13.5 // indirect
cosmossdk.io/x/upgrade v0.1.4 // indirect
Expand Down Expand Up @@ -79,7 +80,6 @@ require (
github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06 // indirect
github.com/cometbft/cometbft-db v0.11.0 // indirect
github.com/cosmos/btcutil v1.0.5 // indirect
github.com/cosmos/cosmos-db v1.0.2 // indirect
github.com/cosmos/go-bip39 v1.0.0 // indirect
github.com/cosmos/gogogateway v1.2.0 // indirect
github.com/cosmos/iavl v1.2.0 // indirect
Expand Down
40 changes: 39 additions & 1 deletion keeper/keeper.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,20 +4,36 @@ import (
"context"
"fmt"

"cosmossdk.io/collections"
"cosmossdk.io/core/address"
"cosmossdk.io/core/event"
"cosmossdk.io/core/store"
"cosmossdk.io/math"
sdk "github.com/cosmos/cosmos-sdk/types"
transfertypes "github.com/cosmos/ibc-go/v8/modules/apps/transfer/types"
"github.com/ondoprotocol/usdy-noble/v2/types"
"github.com/ondoprotocol/usdy-noble/v2/types/blocklist"
)

type Keeper struct {
Denom string

schema collections.Schema
storeService store.KVStoreService
eventService event.Service

Paused collections.Item[bool]
Owner collections.Item[string]
PendingOwner collections.Item[string]
Burners collections.Map[string, math.Int]
Minters collections.Map[string, math.Int]
Pausers collections.Map[string, []byte]
BlockedChannels collections.Map[string, []byte]

BlocklistOwner collections.Item[string]
BlocklistPendingOwner collections.Item[string]
BlockedAddresses collections.Map[[]byte, []byte]

addressCodec address.Codec
bankKeeper types.BankKeeper
}
Expand All @@ -29,15 +45,37 @@ func NewKeeper(
addressCodec address.Codec,
bankKeeper types.BankKeeper,
) *Keeper {
return &Keeper{
builder := collections.NewSchemaBuilder(storeService)

keeper := &Keeper{
Denom: denom,

storeService: storeService,
eventService: eventService,

Paused: collections.NewItem(builder, types.PausedKey, "paused", collections.BoolValue),
Owner: collections.NewItem(builder, types.OwnerKey, "owner", collections.StringValue),
PendingOwner: collections.NewItem(builder, types.PendingOwnerKey, "pending_owner", collections.StringValue),
Burners: collections.NewMap(builder, types.BurnerPrefix, "burners", collections.StringKey, sdk.IntValue),
Minters: collections.NewMap(builder, types.MinterPrefix, "minters", collections.StringKey, sdk.IntValue),
Pausers: collections.NewMap(builder, types.PauserPrefix, "pausers", collections.StringKey, collections.BytesValue),
BlockedChannels: collections.NewMap(builder, types.BlockedChannelPrefix, "blocked_channels", collections.StringKey, collections.BytesValue),

BlocklistOwner: collections.NewItem(builder, blocklist.OwnerKey, "blocklist_owner", collections.StringValue),
BlocklistPendingOwner: collections.NewItem(builder, blocklist.PendingOwnerKey, "blocklist_pending_owner", collections.StringValue),
BlockedAddresses: collections.NewMap(builder, blocklist.BlockedAddressPrefix, "blocked_address", collections.BytesKey, collections.BytesValue),

addressCodec: addressCodec,
bankKeeper: bankKeeper,
}

schema, err := builder.Build()
if err != nil {
panic(err)
}

keeper.schema = schema
return keeper
}

// SetBankKeeper overwrites the bank keeper used in this module.
Expand Down
71 changes: 47 additions & 24 deletions keeper/keeper_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,11 @@ import (
"testing"

"cosmossdk.io/math"
"github.com/cosmos/cosmos-sdk/codec/address"
"github.com/cosmos/cosmos-sdk/runtime"
sdk "github.com/cosmos/cosmos-sdk/types"
transfertypes "github.com/cosmos/ibc-go/v8/modules/apps/transfer/types"
"github.com/ondoprotocol/usdy-noble/v2/keeper"
"github.com/ondoprotocol/usdy-noble/v2/types"
"github.com/ondoprotocol/usdy-noble/v2/utils"
"github.com/ondoprotocol/usdy-noble/v2/utils/mocks"
Expand All @@ -15,8 +18,8 @@ import (

func TestSendRestrictionBurn(t *testing.T) {
user := utils.TestAccount()
keeper, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(keeper.Denom, ONE))
k, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(k.Denom, ONE))

testCases := []struct {
name string
Expand Down Expand Up @@ -53,16 +56,16 @@ func TestSendRestrictionBurn(t *testing.T) {
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
// ARRANGE: Set paused state.
keeper.SetPaused(ctx, testCase.paused)
require.NoError(t, k.SetPaused(ctx, testCase.paused))
// ARRANGE: Set blocked state.
if testCase.blocked {
keeper.SetBlockedAddress(ctx, user.Bytes)
require.NoError(t, k.SetBlockedAddress(ctx, user.Bytes))
} else {
keeper.DeleteBlockedAddress(ctx, user.Bytes)
require.NoError(t, k.DeleteBlockedAddress(ctx, user.Bytes))
}

// ACT: Attempt to burn.
_, err := keeper.SendRestrictionFn(ctx, user.Bytes, types.ModuleAddress, coins)
_, err := k.SendRestrictionFn(ctx, user.Bytes, types.ModuleAddress, coins)

// ASSERT: Send restriction correctly handled test case.
if testCase.err != nil {
Expand All @@ -76,8 +79,8 @@ func TestSendRestrictionBurn(t *testing.T) {

func TestSendRestrictionMint(t *testing.T) {
user := utils.TestAccount()
keeper, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(keeper.Denom, ONE))
k, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(k.Denom, ONE))

testCases := []struct {
name string
Expand Down Expand Up @@ -114,16 +117,16 @@ func TestSendRestrictionMint(t *testing.T) {
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
// ARRANGE: Set paused state.
keeper.SetPaused(ctx, testCase.paused)
require.NoError(t, k.SetPaused(ctx, testCase.paused))
// ARRANGE: Set blocked state.
if testCase.blocked {
keeper.SetBlockedAddress(ctx, user.Bytes)
require.NoError(t, k.SetBlockedAddress(ctx, user.Bytes))
} else {
keeper.DeleteBlockedAddress(ctx, user.Bytes)
require.NoError(t, k.DeleteBlockedAddress(ctx, user.Bytes))
}

// ACT: Attempt to mint.
_, err := keeper.SendRestrictionFn(ctx, types.ModuleAddress, user.Bytes, coins)
_, err := k.SendRestrictionFn(ctx, types.ModuleAddress, user.Bytes, coins)

// ASSERT: Send restriction correctly handled test case.
if testCase.err != nil {
Expand All @@ -137,8 +140,8 @@ func TestSendRestrictionMint(t *testing.T) {

func TestSendRestrictionTransfer(t *testing.T) {
alice, bob := utils.TestAccount(), utils.TestAccount()
keeper, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(keeper.Denom, ONE))
k, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(k.Denom, ONE))

testCases := []struct {
name string
Expand Down Expand Up @@ -225,22 +228,22 @@ func TestSendRestrictionTransfer(t *testing.T) {
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
// ARRANGE: Set paused state.
keeper.SetPaused(ctx, testCase.paused)
require.NoError(t, k.SetPaused(ctx, testCase.paused))
// ARRANGE: Set sender blocked state.
if testCase.senderBlocked {
keeper.SetBlockedAddress(ctx, alice.Bytes)
require.NoError(t, k.SetBlockedAddress(ctx, alice.Bytes))
} else {
keeper.DeleteBlockedAddress(ctx, alice.Bytes)
require.NoError(t, k.DeleteBlockedAddress(ctx, alice.Bytes))
}
// ARRANGE: Set recipient blocked state.
if testCase.recipientBlocked {
keeper.SetBlockedAddress(ctx, bob.Bytes)
require.NoError(t, k.SetBlockedAddress(ctx, bob.Bytes))
} else {
keeper.DeleteBlockedAddress(ctx, bob.Bytes)
require.NoError(t, k.DeleteBlockedAddress(ctx, bob.Bytes))
}

// ACT: Attempt to transfer.
_, err := keeper.SendRestrictionFn(ctx, alice.Bytes, bob.Bytes, testCase.coins)
_, err := k.SendRestrictionFn(ctx, alice.Bytes, bob.Bytes, testCase.coins)

// ASSERT: Send restriction correctly handled test case.
if testCase.err != nil {
Expand All @@ -254,17 +257,37 @@ func TestSendRestrictionTransfer(t *testing.T) {

func TestSendRestrictionIBCTransfer(t *testing.T) {
user := utils.TestAccount()
keeper, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(keeper.Denom, ONE))
k, ctx := mocks.AuraKeeper()
coins := sdk.NewCoins(sdk.NewCoin(k.Denom, ONE))

// ARRANGE: Set a blocked channel in state.
keeper.SetBlockedChannel(ctx, "channel-0")
require.NoError(t, k.SetBlockedChannel(ctx, "channel-0"))
escrow := transfertypes.GetEscrowAddress(transfertypes.PortID, "channel-0")

// ACT: Attempt to transfer from user to escrow account.
// This is to mimic the underlying transfer that occurs when using IBC.
_, err := keeper.SendRestrictionFn(ctx, user.Bytes, escrow, coins)
_, err := k.SendRestrictionFn(ctx, user.Bytes, escrow, coins)

// ASSERT: The action should've failed due to blocked channel.
require.ErrorContains(t, err, "transfers are blocked")
}

func TestNewKeeper(t *testing.T) {
// ARRANGE: Set the PausedKey to an already existing key
types.PausedKey = types.OwnerKey

// ACT: Attempt to create a new Keeper with overlapping prefixes
require.Panics(t, func() {
keeper.NewKeeper(
"ausdy",
mocks.FailingStore(mocks.Set, nil),
runtime.ProvideEventService(),
address.NewBech32Codec("noble"),
mocks.BankKeeper{},
)
})
// ASSERT: The function should've panicked.

// ARRANGE: Restore the original PausedKey
types.PausedKey = []byte("paused")
}
Loading

0 comments on commit 099f728

Please sign in to comment.