diff --git a/tools/test-framework/src/chain/cli/wasm/query.rs b/tools/test-framework/src/chain/cli/wasm/query.rs
index 14a1f042be..fb2835c54b 100644
--- a/tools/test-framework/src/chain/cli/wasm/query.rs
+++ b/tools/test-framework/src/chain/cli/wasm/query.rs
@@ -46,6 +46,7 @@ pub fn query_wasm_list_contracts_by_code(
     command_path: &str,
     home_path: &str,
     rpc_listen_address: &str,
+    code_id: &str,
 ) -> Result<String, Error> {
     let exec_output = simple_exec(
         chain_id,
@@ -60,7 +61,7 @@ pub fn query_wasm_list_contracts_by_code(
             "query",
             "wasm",
             "list-contract-by-code",
-            "1",
+            code_id,
             "--output",
             "json",
         ],
diff --git a/tools/test-framework/src/chain/ext/async_icq.rs b/tools/test-framework/src/chain/ext/async_icq.rs
index 2e3c846e08..88dd7f7aa6 100644
--- a/tools/test-framework/src/chain/ext/async_icq.rs
+++ b/tools/test-framework/src/chain/ext/async_icq.rs
@@ -37,6 +37,7 @@ impl<'a, Chain: Send> AsyncIcqMethodsExt<Chain> for MonoTagged<Chain, &'a ChainD
             &driver.command_path,
             &driver.home_path,
             &driver.rpc_listen_address(),
+            &wasm_code,
         )?;
 
         update_oracle(