[TRA-418] Add upgrade handler and test for x/revshare module state in…

…itialization (#1783)

Signed-off-by: Shrenuj Bansal <shrenuj@dydx.exchange>

protocol/app/upgrades.go

@@ -31,6 +31,8 @@ func (app *App) setupUpgradeHandlers() {
 			app.ModuleManager,
 			app.configurator,
 			app.ClobKeeper,
+			app.RevShareKeeper,
+			app.PricesKeeper,
 		),
 	)
 }

protocol/app/upgrades/v6.0.0/upgrade.go

@@ -7,11 +7,14 @@ import (
 	upgradetypes "cosmossdk.io/x/upgrade/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	indexerevents "github.com/dydxprotocol/v4-chain/protocol/indexer/events"
 	"github.com/dydxprotocol/v4-chain/protocol/indexer/indexer_manager"
 	indexershared "github.com/dydxprotocol/v4-chain/protocol/indexer/shared"
 	"github.com/dydxprotocol/v4-chain/protocol/lib"
 	clobtypes "github.com/dydxprotocol/v4-chain/protocol/x/clob/types"
+	pricetypes "github.com/dydxprotocol/v4-chain/protocol/x/prices/types"
+	revsharetypes "github.com/dydxprotocol/v4-chain/protocol/x/revshare/types"
 )
 
 func removeStatefulFOKOrders(ctx sdk.Context, k clobtypes.ClobKeeper) {
@@ -39,10 +42,38 @@ func removeStatefulFOKOrders(ctx sdk.Context, k clobtypes.ClobKeeper) {
 	}
 }
 
+func initRevShareModuleState(
+	ctx sdk.Context,
+	revShareKeeper revsharetypes.RevShareKeeper,
+	priceKeeper pricetypes.PricesKeeper,
+) {
+	// Initialize the rev share module state.
+	params := revsharetypes.MarketMapperRevenueShareParams{
+		Address:         authtypes.NewModuleAddress(authtypes.FeeCollectorName).String(),
+		RevenueSharePpm: 0,
+		ValidDays:       0,
+	}
+	err := revShareKeeper.SetMarketMapperRevenueShareParams(ctx, params)
+	if err != nil {
+		panic(fmt.Sprintf("failed to set market mapper revenue share params: %s", err))
+	}
+
+	// Initialize the rev share details for all existing markets.
+	markets := priceKeeper.GetAllMarketParams(ctx)
+	for _, market := range markets {
+		revShareDetails := revsharetypes.MarketMapperRevShareDetails{
+			ExpirationTs: 0,
+		}
+		revShareKeeper.SetMarketMapperRevShareDetails(ctx, market.Id, revShareDetails)
+	}
+}
+
 func CreateUpgradeHandler(
 	mm *module.Manager,
 	configurator module.Configurator,
 	clobKeeper clobtypes.ClobKeeper,
+	revShareKeeper revsharetypes.RevShareKeeper,
+	priceKeeper pricetypes.PricesKeeper,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		sdkCtx := lib.UnwrapSDKContext(ctx, "app/upgrades")
@@ -51,6 +82,9 @@ func CreateUpgradeHandler(
 		// Remove all stateful FOK orders from state.
 		removeStatefulFOKOrders(sdkCtx, clobKeeper)
 
+		// Initialize the rev share module state.
+		initRevShareModuleState(sdkCtx, revShareKeeper, priceKeeper)
+
 		sdkCtx.Logger().Info("Successfully removed stateful orders from state")
 
 		return mm.RunMigrations(ctx, configurator, vm)

protocol/app/upgrades/v6.0.0/upgrade_container_test.go

@@ -6,6 +6,12 @@ import (
 	"testing"
 	"time"
 
+	pricestypes "github.com/dydxprotocol/v4-chain/protocol/x/prices/types"
+
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	"github.com/cosmos/gogoproto/proto"
+	revsharetypes "github.com/dydxprotocol/v4-chain/protocol/x/revshare/types"
+
 	v_6_0_0 "github.com/dydxprotocol/v4-chain/protocol/app/upgrades/v6.0.0"
 	"github.com/dydxprotocol/v4-chain/protocol/testing/containertest"
 	"github.com/dydxprotocol/v4-chain/protocol/testutil/constants"
@@ -50,6 +56,7 @@ func preUpgradeChecks(node *containertest.Node, t *testing.T) {
 func postUpgradeChecks(node *containertest.Node, t *testing.T) {
 	// Add test for your upgrade handler logic below
 	postUpgradeStatefulOrderCheck(node, t)
+	postUpgradeMarketMapperRevShareChecks(node, t)
 }
 
 func placeOrders(node *containertest.Node, t *testing.T) {
@@ -192,3 +199,50 @@ func postUpgradeStatefulOrderCheck(node *containertest.Node, t *testing.T) {
 	)
 	require.ErrorIs(t, err, status.Error(codes.NotFound, "not found"))
 }
+
+func postUpgradeMarketMapperRevShareChecks(node *containertest.Node, t *testing.T) {
+	// Check that all rev share params are set to the default value
+	resp, err := containertest.Query(
+		node,
+		revsharetypes.NewQueryClient,
+		revsharetypes.QueryClient.MarketMapperRevenueShareParams,
+		&revsharetypes.QueryMarketMapperRevenueShareParams{},
+	)
+	require.NoError(t, err)
+
+	params := revsharetypes.QueryMarketMapperRevenueShareParamsResponse{}
+	err = proto.UnmarshalText(resp.String(), &params)
+	require.NoError(t, err)
+	require.Equal(t, params.Params.Address, authtypes.NewModuleAddress(authtypes.FeeCollectorName).String())
+	require.Equal(t, params.Params.RevenueSharePpm, uint32(0))
+	require.Equal(t, params.Params.ValidDays, uint32(0))
+
+	// Get all markets list
+	marketParams := pricestypes.QueryAllMarketParamsResponse{}
+	resp, err = containertest.Query(
+		node,
+		pricestypes.NewQueryClient,
+		pricestypes.QueryClient.AllMarketParams,
+		&pricestypes.QueryAllMarketParamsRequest{},
+	)
+	require.NoError(t, err)
+	err = proto.UnmarshalText(resp.String(), &marketParams)
+	require.NoError(t, err)
+
+	// Check that all rev share details are set to the default value
+	for _, market := range marketParams.MarketParams {
+		revShareDetails := revsharetypes.QueryMarketMapperRevShareDetailsResponse{}
+		resp, err := containertest.Query(
+			node,
+			revsharetypes.NewQueryClient,
+			revsharetypes.QueryClient.MarketMapperRevShareDetails,
+			&revsharetypes.QueryMarketMapperRevShareDetails{
+				MarketId: market.Id,
+			},
+		)
+		require.NoError(t, err)
+		err = proto.UnmarshalText(resp.String(), &revShareDetails)
+		require.NoError(t, err)
+		require.Equal(t, revShareDetails.Details.ExpirationTs, uint64(0))
+	}
+}

protocol/x/revshare/types/types.go

@@ -0,0 +1,26 @@
+package types
+
+import sdk "github.com/cosmos/cosmos-sdk/types"
+
+type RevShareKeeper interface {
+	// MarketMapperRevenueShareParams
+	SetMarketMapperRevenueShareParams(
+		ctx sdk.Context,
+		params MarketMapperRevenueShareParams,
+	) (err error)
+	GetMarketMapperRevenueShareParams(
+		ctx sdk.Context,
+	) (params MarketMapperRevenueShareParams)
+
+	// MarketMapperRevShareDetails
+	SetMarketMapperRevShareDetails(
+		ctx sdk.Context,
+		marketId uint32,
+		params MarketMapperRevShareDetails,
+	)
+	GetMarketMapperRevShareDetails(
+		ctx sdk.Context,
+		marketId uint32,
+	) (params MarketMapperRevShareDetails, err error)
+	CreateNewMarketRevShare(ctx sdk.Context, marketId uint32)
+}