From 68cf67730ac4bb3e2ceae087d507df2f1cd87148 Mon Sep 17 00:00:00 2001
From: Xiaoyu Chen <codchen03@gmail.com>
Date: Mon, 20 Jun 2022 14:31:59 -0700
Subject: [PATCH] fix migration and add tests

---
 x/dex/migrations/v3_to_v4.go      |  4 +--
 x/dex/migrations/v3_to_v4_test.go | 48 +++++++++++++++++++++++++++++++
 x/dex/module.go                   |  2 +-
 3 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 x/dex/migrations/v3_to_v4_test.go

diff --git a/x/dex/migrations/v3_to_v4.go b/x/dex/migrations/v3_to_v4.go
index 19976d7115..b738c6934c 100644
--- a/x/dex/migrations/v3_to_v4.go
+++ b/x/dex/migrations/v3_to_v4.go
@@ -1,13 +1,13 @@
 package migrations
 
 import (
-	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
 	"github.com/sei-protocol/sei-chain/x/dex/types"
 )
 
-func PriceSnapshotUpdate(ctx sdk.Context, storeKey sdk.StoreKey, cdc codec.BinaryCodec, paramStore paramtypes.Subspace) error {
+func PriceSnapshotUpdate(ctx sdk.Context, paramStore paramtypes.Subspace) error {
+	migratePriceSnapshotParam(ctx, paramStore)
 	return nil
 }
 
diff --git a/x/dex/migrations/v3_to_v4_test.go b/x/dex/migrations/v3_to_v4_test.go
new file mode 100644
index 0000000000..18712ba897
--- /dev/null
+++ b/x/dex/migrations/v3_to_v4_test.go
@@ -0,0 +1,48 @@
+package migrations_test
+
+import (
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	"github.com/cosmos/cosmos-sdk/store"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	typesparams "github.com/cosmos/cosmos-sdk/x/params/types"
+	"github.com/sei-protocol/sei-chain/x/dex/migrations"
+	"github.com/sei-protocol/sei-chain/x/dex/types"
+	"github.com/stretchr/testify/require"
+	"github.com/tendermint/tendermint/libs/log"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+	tmdb "github.com/tendermint/tm-db"
+)
+
+func TestMigrate3to4(t *testing.T) {
+	storeKey := sdk.NewKVStoreKey(types.StoreKey)
+	memStoreKey := storetypes.NewMemoryStoreKey(types.MemStoreKey)
+
+	db := tmdb.NewMemDB()
+	stateStore := store.NewCommitMultiStore(db)
+	stateStore.MountStoreWithDB(storeKey, sdk.StoreTypeIAVL, db)
+	stateStore.MountStoreWithDB(memStoreKey, sdk.StoreTypeMemory, nil)
+	require.NoError(t, stateStore.LoadLatestVersion())
+
+	registry := codectypes.NewInterfaceRegistry()
+	cdc := codec.NewProtoCodec(registry)
+
+	paramsSubspace := typesparams.NewSubspace(cdc,
+		types.Amino,
+		storeKey,
+		memStoreKey,
+		"DexParams",
+	)
+	ctx := sdk.NewContext(stateStore, tmproto.Header{}, false, log.NewNopLogger())
+	if !paramsSubspace.HasKeyTable() {
+		paramsSubspace = paramsSubspace.WithKeyTable(types.ParamKeyTable())
+	}
+	migrations.PriceSnapshotUpdate(ctx, paramsSubspace)
+
+	params := types.Params{}
+	paramsSubspace.GetParamSet(ctx, &params)
+	require.Equal(t, uint64(types.DefaultPriceSnapshotRetention), params.PriceSnapshotRetention)
+}
diff --git a/x/dex/module.go b/x/dex/module.go
index f7e4a7c95e..3da9f00a12 100644
--- a/x/dex/module.go
+++ b/x/dex/module.go
@@ -159,7 +159,7 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 		return migrations.DataTypeUpdate(ctx, am.keeper.GetStoreKey(), am.keeper.Cdc)
 	})
 	cfg.RegisterMigration(types.ModuleName, 3, func(ctx sdk.Context) error {
-		return migrations.PriceSnapshotUpdate(ctx, am.keeper.GetStoreKey(), am.keeper.Cdc, am.keeper.Paramstore)
+		return migrations.PriceSnapshotUpdate(ctx, am.keeper.Paramstore)
 	})
 }