From 0e6d7243e9304db3371d18a292d8bd6bc8a94192 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Thu, 6 Aug 2020 12:30:07 +0200
Subject: [PATCH] Support gaps in query code list

---
 x/wasm/internal/keeper/querier.go      | 12 ++-----
 x/wasm/internal/keeper/querier_test.go | 50 ++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/x/wasm/internal/keeper/querier.go b/x/wasm/internal/keeper/querier.go
index 4fcc2a027c..b972cc6aae 100644
--- a/x/wasm/internal/keeper/querier.go
+++ b/x/wasm/internal/keeper/querier.go
@@ -205,14 +205,7 @@ type ListCodeResponse struct {
 
 func queryCodeList(ctx sdk.Context, keeper Keeper) ([]byte, error) {
 	var info []ListCodeResponse
-
-	var i uint64
-	for true {
-		i++
-		res := keeper.GetCodeInfo(ctx, i)
-		if res == nil {
-			break
-		}
+	keeper.IterateCodeInfos(ctx, func(i uint64, res types.CodeInfo) bool {
 		info = append(info, ListCodeResponse{
 			ID:       i,
 			Creator:  res.Creator,
@@ -220,7 +213,8 @@ func queryCodeList(ctx sdk.Context, keeper Keeper) ([]byte, error) {
 			Source:   res.Source,
 			Builder:  res.Builder,
 		})
-	}
+		return false
+	})
 
 	bz, err := json.MarshalIndent(info, "", "  ")
 	if err != nil {
diff --git a/x/wasm/internal/keeper/querier_test.go b/x/wasm/internal/keeper/querier_test.go
index 599bd39d6e..0a37b8931d 100644
--- a/x/wasm/internal/keeper/querier_test.go
+++ b/x/wasm/internal/keeper/querier_test.go
@@ -318,3 +318,53 @@ func TestQueryContractHistory(t *testing.T) {
 		})
 	}
 }
+
+func TestQueryCodeList(t *testing.T) {
+	wasmCode, err := ioutil.ReadFile("./testdata/contract.wasm")
+	require.NoError(t, err)
+
+	specs := map[string]struct {
+		codeIDs []uint64
+	}{
+		"none": {},
+		"no gaps": {
+			codeIDs: []uint64{1, 2, 3},
+		},
+		"with gaps": {
+			codeIDs: []uint64{2, 4, 6},
+		},
+	}
+
+	for msg, spec := range specs {
+		t.Run(msg, func(t *testing.T) {
+			tempDir, err := ioutil.TempDir("", "wasm")
+			require.NoError(t, err)
+			defer os.RemoveAll(tempDir)
+			ctx, keepers := CreateTestInput(t, false, tempDir, SupportedFeatures, nil, nil)
+			keeper := keepers.WasmKeeper
+
+			for _, codeID := range spec.codeIDs {
+				require.NoError(t, keeper.importCode(ctx, codeID,
+					types.CodeInfoFixture(types.WithSHA256CodeHash(wasmCode)),
+					wasmCode),
+				)
+			}
+			q := NewQuerier(keeper)
+			// when
+			query := []string{QueryListCode}
+			data := abci.RequestQuery{}
+			resData, err := q(ctx, query, data)
+
+			// then
+			require.NoError(t, err)
+
+			var got []map[string]interface{}
+			err = json.Unmarshal(resData, &got)
+			require.NoError(t, err)
+			require.Len(t, got, len(spec.codeIDs))
+			for i, exp := range spec.codeIDs {
+				assert.EqualValues(t, exp, got[i]["id"])
+			}
+		})
+	}
+}