From 8c7967e5b48a48a459fe783356b175afedb307d7 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Tue, 27 Apr 2021 09:04:58 +0200
Subject: [PATCH] Review feedback

---
 x/wasm/keeper/keeper_test.go  | 27 +++++++++++++++++++++++++++
 x/wasm/keeper/querier_test.go |  2 +-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/x/wasm/keeper/keeper_test.go b/x/wasm/keeper/keeper_test.go
index 4b5607f7bd..fc88de3b3e 100644
--- a/x/wasm/keeper/keeper_test.go
+++ b/x/wasm/keeper/keeper_test.go
@@ -1096,6 +1096,33 @@ func TestIterateContractsByCode(t *testing.T) {
 	}
 }
 
+func TestIterateContractsByCodeWithMigration(t *testing.T) {
+	// mock migration so that it does not fail when migrate example1 to example2.codeID
+	mockWasmVM := wasmtesting.MockWasmer{MigrateFn: func(codeID wasmvm.Checksum, env wasmvmtypes.Env, migrateMsg []byte, store wasmvm.KVStore, goapi wasmvm.GoAPI, querier wasmvm.Querier, gasMeter wasmvm.GasMeter, gasLimit uint64) (*wasmvmtypes.Response, uint64, error) {
+		return &wasmvmtypes.Response{}, 1, nil
+	}}
+	wasmtesting.MakeInstantiable(&mockWasmVM)
+	ctx, keepers := CreateTestInput(t, false, SupportedFeatures, WithWasmEngine(&mockWasmVM))
+	k, c := keepers.WasmKeeper, keepers.ContractKeeper
+	example1 := InstantiateHackatomExampleContract(t, ctx, keepers)
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
+	example2 := InstantiateIBCReflectContract(t, ctx, keepers)
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
+	_, err := c.Migrate(ctx, example1.Contract, example1.CreatorAddr, example2.CodeID, []byte("{}"))
+	require.NoError(t, err)
+
+	// when
+	var gotAddr []sdk.AccAddress
+	k.IterateContractsByCode(ctx, example2.CodeID, func(address sdk.AccAddress) bool {
+		gotAddr = append(gotAddr, address)
+		return false
+	})
+
+	// then
+	exp := []sdk.AccAddress{example2.Contract, example1.Contract}
+	assert.Equal(t, exp, gotAddr)
+}
+
 type sudoMsg struct {
 	// This is a tongue-in-check demo command. This is not the intended purpose of Sudo.
 	// Here we show that some priviledged Go module can make a call that should never be exposed
diff --git a/x/wasm/keeper/querier_test.go b/x/wasm/keeper/querier_test.go
index cdb7098539..ee989014a1 100644
--- a/x/wasm/keeper/querier_test.go
+++ b/x/wasm/keeper/querier_test.go
@@ -289,7 +289,7 @@ func TestQueryContractListByCodeOrdering(t *testing.T) {
 			ctx = setBlock(ctx, h)
 			h++
 		}
-		_, _, err = keepers.ContractKeeper.Instantiate(ctx, codeID, creator, nil, initMsgBz, fmt.Sprintf("contractAddr %d", i), topUp)
+		_, _, err = keepers.ContractKeeper.Instantiate(ctx, codeID, creator, nil, initMsgBz, fmt.Sprintf("contract %d", i), topUp)
 		require.NoError(t, err)
 	}