From ac61de2751080116f6a54235f84046a641f45fb0 Mon Sep 17 00:00:00 2001
From: affan <affan@dydx.exchange>
Date: Tue, 24 Sep 2024 13:25:12 -0400
Subject: [PATCH] upgrade handler

---
 protocol/app/upgrades.go                      |  1 +
 protocol/app/upgrades/v7.0.0/upgrade.go       | 10 +++++++++
 .../upgrades/v7.0.0/upgrade_container_test.go | 21 +++++++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/protocol/app/upgrades.go b/protocol/app/upgrades.go
index df4bf3eb38..1abb1b3764 100644
--- a/protocol/app/upgrades.go
+++ b/protocol/app/upgrades.go
@@ -32,6 +32,7 @@ func (app *App) setupUpgradeHandlers() {
 			app.configurator,
 			app.PricesKeeper,
 			app.VaultKeeper,
+			app.AffiliatesKeeper,
 		),
 	)
 }
diff --git a/protocol/app/upgrades/v7.0.0/upgrade.go b/protocol/app/upgrades/v7.0.0/upgrade.go
index 14a2b670b1..d566b94727 100644
--- a/protocol/app/upgrades/v7.0.0/upgrade.go
+++ b/protocol/app/upgrades/v7.0.0/upgrade.go
@@ -9,6 +9,8 @@ import (
 	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/dydxprotocol/v4-chain/protocol/lib"
 	"github.com/dydxprotocol/v4-chain/protocol/lib/slinky"
+	affiliateskeeper "github.com/dydxprotocol/v4-chain/protocol/x/affiliates/keeper"
+	affiliatetypes "github.com/dydxprotocol/v4-chain/protocol/x/affiliates/types"
 	pricestypes "github.com/dydxprotocol/v4-chain/protocol/x/prices/types"
 	vaultkeeper "github.com/dydxprotocol/v4-chain/protocol/x/vault/keeper"
 	vaulttypes "github.com/dydxprotocol/v4-chain/protocol/x/vault/types"
@@ -55,11 +57,16 @@ func migrateVaultQuotingParamsToVaultParams(ctx sdk.Context, k vaultkeeper.Keepe
 	}
 }
 
+func setDefaultTiersForAffiliates(ctx sdk.Context, k affiliateskeeper.Keeper) {
+	k.UpdateAffiliateTiers(ctx, affiliatetypes.DefaultAffiliateTiers)
+}
+
 func CreateUpgradeHandler(
 	mm *module.Manager,
 	configurator module.Configurator,
 	pricesKeeper pricestypes.PricesKeeper,
 	vaultKeeper vaultkeeper.Keeper,
+	affiliatesKeeper affiliateskeeper.Keeper,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		sdkCtx := lib.UnwrapSDKContext(ctx, "app/upgrades")
@@ -71,6 +78,9 @@ func CreateUpgradeHandler(
 		// Migrate vault quoting params to vault params.
 		migrateVaultQuotingParamsToVaultParams(sdkCtx, vaultKeeper)
 
+		// Set default tiers for affiliates.
+		setDefaultTiersForAffiliates(sdkCtx, affiliatesKeeper)
+
 		return mm.RunMigrations(ctx, configurator, vm)
 	}
 }
diff --git a/protocol/app/upgrades/v7.0.0/upgrade_container_test.go b/protocol/app/upgrades/v7.0.0/upgrade_container_test.go
index e92ee609ca..f4d9e556a2 100644
--- a/protocol/app/upgrades/v7.0.0/upgrade_container_test.go
+++ b/protocol/app/upgrades/v7.0.0/upgrade_container_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/dydxprotocol/v4-chain/protocol/dtypes"
 	"github.com/dydxprotocol/v4-chain/protocol/testing/containertest"
 	"github.com/dydxprotocol/v4-chain/protocol/testutil/constants"
+	affiliatestypes "github.com/dydxprotocol/v4-chain/protocol/x/affiliates/types"
 	vaulttypes "github.com/dydxprotocol/v4-chain/protocol/x/vault/types"
 	"github.com/stretchr/testify/require"
 )
@@ -48,6 +49,9 @@ func preUpgradeChecks(node *containertest.Node, t *testing.T) {
 func postUpgradeChecks(node *containertest.Node, t *testing.T) {
 	// Add test for your upgrade handler logic below
 	postUpgradeVaultParamsCheck(node, t)
+
+	// Check that the affiliates module has been initialized with the default tiers.
+	postUpgradeAffiliatesModuleTiersCheck(node, t)
 }
 
 func postUpgradeVaultParamsCheck(node *containertest.Node, t *testing.T) {
@@ -96,3 +100,20 @@ func checkVaultParams(
 	require.Equal(t, expectedStatus, vaultParamsResp.VaultParams.Status)
 	require.Equal(t, expectedQuotingParams, vaultParamsResp.VaultParams.QuotingParams)
 }
+
+func postUpgradeAffiliatesModuleTiersCheck(node *containertest.Node, t *testing.T) {
+
+	resp, err := containertest.Query(
+		node,
+		affiliatestypes.NewQueryClient,
+		affiliatestypes.QueryClient.AllAffiliateTiers,
+		&affiliatestypes.AllAffiliateTiersRequest{},
+	)
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	affiliateTiersResp := affiliatestypes.AllAffiliateTiersResponse{}
+	err = proto.UnmarshalText(resp.String(), &affiliateTiersResp)
+	require.NoError(t, err)
+	require.Equal(t, affiliatestypes.DefaultAffiliateTiers, affiliateTiersResp.Tiers)
+}