refactor: utilize collections for state management

genesis.go

@@ -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)
+		}
 	}
 }
 

go.mod

@@ -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
@@ -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
@@ -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

keeper/keeper.go

@@ -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
 }
@@ -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.

keeper/keeper_test.go

@@ -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"
@@ -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
@@ -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 {
@@ -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
@@ -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 {
@@ -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
@@ -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 {
@@ -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")
+}