From 669d18a24b8d811efa30999edd5564accdac7b4a Mon Sep 17 00:00:00 2001
From: Yair Bakalchuk <yair@starkware.co>
Date: Thu, 4 Jul 2024 15:09:35 +0300
Subject: [PATCH] refactor(gateway): prepare run_validate for mocking

---
 crates/gateway/src/gateway.rs                 |  3 +-
 .../src/stateful_transaction_validator.rs     | 54 ++++++++++++++-----
 .../stateful_transaction_validator_test.rs    | 10 ++--
 3 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/crates/gateway/src/gateway.rs b/crates/gateway/src/gateway.rs
index 352a80d51..1e809eb24 100644
--- a/crates/gateway/src/gateway.rs
+++ b/crates/gateway/src/gateway.rs
@@ -124,8 +124,9 @@ fn process_tx(
     };
 
     // TODO(Yael, 19/5/2024): pass the relevant deploy_account_hash.
+    let blockifier_validator = stateful_tx_validator.prepare_validate(state_reader_factory)?;
     let tx_hash =
-        stateful_tx_validator.run_validate(state_reader_factory, &tx, optional_class_info, None)?;
+        stateful_tx_validator.run_validate(&tx, optional_class_info, None, blockifier_validator)?;
 
     // TODO(Arni): Add the Sierra and the Casm to the mempool input.
     Ok(MempoolInput {
diff --git a/crates/gateway/src/stateful_transaction_validator.rs b/crates/gateway/src/stateful_transaction_validator.rs
index c45a28878..316a68cd3 100644
--- a/crates/gateway/src/stateful_transaction_validator.rs
+++ b/crates/gateway/src/stateful_transaction_validator.rs
@@ -4,6 +4,7 @@ use blockifier::bouncer::BouncerConfig;
 use blockifier::context::BlockContext;
 use blockifier::execution::contract_class::ClassInfo;
 use blockifier::state::cached_state::CachedState;
+use blockifier::transaction::account_transaction::AccountTransaction;
 use blockifier::versioned_constants::VersionedConstants;
 use starknet_api::rpc_transaction::RPCTransaction;
 use starknet_api::transaction::TransactionHash;
@@ -21,14 +22,49 @@ pub struct StatefulTransactionValidator {
     pub config: StatefulTransactionValidatorConfig,
 }
 
+pub trait StatefulTransactionValidatorTrait {
+    fn validate(
+        &mut self,
+        account_tx: AccountTransaction,
+        deploy_account_tx_hash: Option<TransactionHash>,
+    ) -> StatefulTransactionValidatorResult<()>;
+}
+
+type ConcreteBlockifierStatefulValidator =
+    BlockifierStatefulValidator<CachedState<Box<dyn MempoolStateReader>>>;
+
+impl StatefulTransactionValidatorTrait for ConcreteBlockifierStatefulValidator {
+    fn validate(
+        &mut self,
+        account_tx: AccountTransaction,
+        deploy_account_tx_hash: Option<TransactionHash>,
+    ) -> StatefulTransactionValidatorResult<()> {
+        Ok(self.perform_validations(account_tx, deploy_account_tx_hash)?)
+    }
+}
+
 impl StatefulTransactionValidator {
-    pub fn run_validate(
+    pub fn run_validate<TStatefulTransactionValidator: StatefulTransactionValidatorTrait>(
         &self,
-        state_reader_factory: &dyn StateReaderFactory,
         external_tx: &RPCTransaction,
         optional_class_info: Option<ClassInfo>,
         deploy_account_tx_hash: Option<TransactionHash>,
+        mut validator: TStatefulTransactionValidator,
     ) -> StatefulTransactionValidatorResult<TransactionHash> {
+        let account_tx = external_tx_to_account_tx(
+            external_tx,
+            optional_class_info,
+            &self.config.chain_info.chain_id,
+        )?;
+        let tx_hash = get_tx_hash(&account_tx);
+        validator.validate(account_tx, deploy_account_tx_hash)?;
+        Ok(tx_hash)
+    }
+
+    pub fn prepare_validate(
+        &self,
+        state_reader_factory: &dyn StateReaderFactory,
+    ) -> StatefulTransactionValidatorResult<ConcreteBlockifierStatefulValidator> {
         // TODO(yael 6/5/2024): consider storing the block_info as part of the
         // StatefulTransactionValidator and update it only once a new block is created.
         let latest_block_info = get_latest_block_info(state_reader_factory)?;
@@ -52,20 +88,12 @@ impl StatefulTransactionValidator {
             &versioned_constants,
         );
 
-        let mut validator = BlockifierStatefulValidator::create(
-            state,
+        Ok(BlockifierStatefulValidator::create(
+            state.into(),
             block_context,
             self.config.max_nonce_for_validation_skip,
             BouncerConfig::max(),
-        );
-        let account_tx = external_tx_to_account_tx(
-            external_tx,
-            optional_class_info,
-            &self.config.chain_info.chain_id,
-        )?;
-        let tx_hash = get_tx_hash(&account_tx);
-        validator.perform_validations(account_tx, deploy_account_tx_hash)?;
-        Ok(tx_hash)
+        ))
     }
 }
 
diff --git a/crates/gateway/src/stateful_transaction_validator_test.rs b/crates/gateway/src/stateful_transaction_validator_test.rs
index 6349f8a1c..1b670863a 100644
--- a/crates/gateway/src/stateful_transaction_validator_test.rs
+++ b/crates/gateway/src/stateful_transaction_validator_test.rs
@@ -85,11 +85,9 @@ fn test_stateful_tx_validator(
         _ => None,
     };
 
-    let result = stateful_validator.run_validate(
-        &state_reader_factory,
-        &external_tx,
-        optional_class_info,
-        None,
-    );
+    let validator = stateful_validator.prepare_validate(&state_reader_factory).unwrap();
+
+    let result =
+        stateful_validator.run_validate(&external_tx, optional_class_info, None, validator);
     assert_eq!(format!("{:?}", result), format!("{:?}", expected_result));
 }