From a64cf6cb0676ac730e93c80363ce48e3dc9394d6 Mon Sep 17 00:00:00 2001
From: AvivYossef-starkware <aviv.y@starkware.co>
Date: Sun, 3 Nov 2024 11:58:14 +0200
Subject: [PATCH] fix(blockifier_reexecution): fix bug in get contract class

---
 .../src/state_reader/test_state_reader.rs        | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs b/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
index 252c4823ef..ec1e49461e 100644
--- a/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
+++ b/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
@@ -21,7 +21,7 @@ use starknet_api::state::StorageKey;
 use starknet_api::transaction::{Transaction, TransactionHash};
 use starknet_core::types::ContractClass as StarknetContractClass;
 use starknet_gateway::config::RpcStateReaderConfig;
-use starknet_gateway::errors::serde_err_to_state_err;
+use starknet_gateway::errors::{serde_err_to_state_err, RPCStateReaderError};
 use starknet_gateway::rpc_objects::{BlockHeader, GetBlockWithTxHashesParams, ResourcePrice};
 use starknet_gateway::rpc_state_reader::RpcStateReader;
 use starknet_types_core::felt::Felt;
@@ -357,10 +357,16 @@ impl ReexecutionStateReader for TestStateReader {
             "block_id": self.rpc_state_reader.block_id,
             "class_hash": class_hash.0.to_string(),
         });
-        let contract_class: StarknetContractClass = serde_json::from_value(
-            self.rpc_state_reader.send_rpc_request("starknet_getClass", params.clone())?,
-        )
-        .map_err(serde_err_to_state_err)?;
+        let raw_contract_class =
+            match self.rpc_state_reader.send_rpc_request("starknet_getClass", params.clone()) {
+                Err(RPCStateReaderError::ClassHashNotFound(_)) => {
+                    return Err(StateError::UndeclaredClassHash(*class_hash));
+                }
+                other_result => other_result,
+            }?;
+
+        let contract_class: StarknetContractClass =
+            serde_json::from_value(raw_contract_class).map_err(serde_err_to_state_err)?;
         // Create a binding to avoid value being dropped.
         let mut dumper_binding = self.contract_class_mapping_dumper.lock().unwrap();
         // If dumper exists, insert the contract class to the mapping.