fix tests

Author: kariy

crates/sync/stage/tests/block.rs

@@ -17,7 +17,7 @@ use katana_primitives::state::StateUpdatesWithClasses;
 use katana_primitives::{felt, ContractAddress, Felt};
 use katana_provider::api::block::{BlockHashProvider, BlockNumberProvider, BlockWriter};
 use katana_provider::test_utils::test_provider;
-use katana_provider::{ProviderError, ProviderResult};
+use katana_provider::{ProviderError, ProviderFactory, ProviderResult};
 use katana_stage::blocks::{BatchBlockDownloader, BlockDownloader, Blocks};
 use katana_stage::{Stage, StageExecutionInput};
 use rstest::rstest;
@@ -119,11 +119,11 @@ impl BlockDownloader for MockBlockDownloader {
     }
 }
 
-/// Mock BlockWriter implementation for testing.
+/// Mock provider implementation for testing.
 ///
 /// Tracks all insert operations and can be configured to return errors.
-#[derive(Clone)]
-struct MockProvider {
+#[derive(Clone, Debug)]
+struct MockInnerProvider {
     /// Stored blocks with their receipts and state updates.
     blocks: Arc<Mutex<Vec<(SealedBlockWithStatus, StateUpdatesWithClasses, Vec<Receipt>)>>>,
     /// Whether to return an error on insert.
@@ -132,40 +132,17 @@ struct MockProvider {
     error_message: Arc<Mutex<String>>,
 }
 
-impl MockProvider {
-    fn new() -> Self {
-        Self {
-            blocks: Arc::new(Mutex::new(Vec::new())),
-            should_fail: Arc::new(Mutex::new(false)),
-            error_message: Arc::new(Mutex::new(String::new())),
-        }
-    }
-
-    /// Add a block directly to the provider's storage.
-    fn with_block(self, block: SealedBlockWithStatus) -> Self {
-        self.blocks.lock().unwrap().push((block, Default::default(), Default::default()));
-        self
-    }
-
-    /// Configure the mock to fail on insert operations.
-    fn with_insert_error(self, error: String) -> Self {
-        *self.should_fail.lock().unwrap() = true;
-        *self.error_message.lock().unwrap() = error;
-        self
-    }
-
-    /// Get the number of blocks stored.
-    fn stored_block_count(&self) -> usize {
-        self.blocks.lock().unwrap().len()
-    }
-
-    /// Get all stored block numbers.
-    fn stored_block_numbers(&self) -> Vec<BlockNumber> {
-        self.blocks.lock().unwrap().iter().map(|(block, _, _)| block.block.header.number).collect()
+impl MockInnerProvider {
+    fn new(
+        blocks: Arc<Mutex<Vec<(SealedBlockWithStatus, StateUpdatesWithClasses, Vec<Receipt>)>>>,
+        should_fail: Arc<Mutex<bool>>,
+        error_message: Arc<Mutex<String>>,
+    ) -> Self {
+        Self { blocks, should_fail, error_message }
     }
 }
 
-impl BlockWriter for MockProvider {
+impl BlockWriter for MockInnerProvider {
     fn insert_block_with_states_and_receipts(
         &self,
         block: SealedBlockWithStatus,
@@ -184,7 +161,7 @@ impl BlockWriter for MockProvider {
     }
 }
 
-impl BlockHashProvider for MockProvider {
+impl BlockHashProvider for MockInnerProvider {
     fn latest_hash(&self) -> ProviderResult<BlockHash> {
         self.blocks
             .lock()
@@ -205,6 +182,85 @@ impl BlockHashProvider for MockProvider {
     }
 }
 
+impl katana_provider::MutableProvider for MockInnerProvider {
+    fn commit(self) -> ProviderResult<()> {
+        Ok(())
+    }
+}
+
+/// Mock ProviderFactory implementation for testing.
+///
+/// Tracks all insert operations and can be configured to return errors.
+#[derive(Clone)]
+struct MockProvider {
+    /// Stored blocks with their receipts and state updates.
+    blocks: Arc<Mutex<Vec<(SealedBlockWithStatus, StateUpdatesWithClasses, Vec<Receipt>)>>>,
+    /// Whether to return an error on insert.
+    should_fail: Arc<Mutex<bool>>,
+    /// Error message to return when should_fail is true.
+    error_message: Arc<Mutex<String>>,
+}
+
+impl std::fmt::Debug for MockProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MockProvider").finish_non_exhaustive()
+    }
+}
+
+impl MockProvider {
+    fn new() -> Self {
+        Self {
+            blocks: Arc::new(Mutex::new(Vec::new())),
+            should_fail: Arc::new(Mutex::new(false)),
+            error_message: Arc::new(Mutex::new(String::new())),
+        }
+    }
+
+    /// Add a block directly to the provider's storage.
+    fn with_block(self, block: SealedBlockWithStatus) -> Self {
+        self.blocks.lock().unwrap().push((block, Default::default(), Default::default()));
+        self
+    }
+
+    /// Configure the mock to fail on insert operations.
+    fn with_insert_error(self, error: String) -> Self {
+        *self.should_fail.lock().unwrap() = true;
+        *self.error_message.lock().unwrap() = error;
+        self
+    }
+
+    /// Get the number of blocks stored.
+    fn stored_block_count(&self) -> usize {
+        self.blocks.lock().unwrap().len()
+    }
+
+    /// Get all stored block numbers.
+    fn stored_block_numbers(&self) -> Vec<BlockNumber> {
+        self.blocks.lock().unwrap().iter().map(|(block, _, _)| block.block.header.number).collect()
+    }
+}
+
+impl katana_provider::ProviderFactory for MockProvider {
+    type Provider = MockInnerProvider;
+    type ProviderMut = MockInnerProvider;
+
+    fn provider(&self) -> Self::Provider {
+        MockInnerProvider::new(
+            Arc::clone(&self.blocks),
+            Arc::clone(&self.should_fail),
+            Arc::clone(&self.error_message),
+        )
+    }
+
+    fn provider_mut(&self) -> Self::ProviderMut {
+        MockInnerProvider::new(
+            Arc::clone(&self.blocks),
+            Arc::clone(&self.should_fail),
+            Arc::clone(&self.error_message),
+        )
+    }
+}
+
 /// Helper function to create a minimal test `SealedBlockWithStatus`.
 ///
 /// Creates a block with the given number and automatically sets the parent hash
@@ -403,13 +459,14 @@ async fn fetch_blocks_from_gateway() {
     let feeder_gateway = SequencerGateway::sepolia();
     let downloader = BatchBlockDownloader::new_gateway(feeder_gateway, 10);
 
-    let mut stage = Blocks::new(&provider, downloader);
+    let mut stage = Blocks::new(provider.clone(), downloader);
 
     let input = StageExecutionInput::new(from_block, to_block);
     stage.execute(&input).await.expect("failed to execute stage");
 
     // check provider storage
-    let block_number = provider.latest_number().expect("failed to get latest block number");
+    let block_number =
+        provider.provider().latest_number().expect("failed to get latest block number");
     assert_eq!(block_number, to_block);
 }
 

crates/sync/stage/tests/trie.rs

@@ -9,63 +9,86 @@ use katana_primitives::Felt;
 use katana_provider::api::block::HeaderProvider;
 use katana_provider::api::state_update::StateUpdateProvider;
 use katana_provider::api::trie::TrieWriter;
-use katana_provider::ProviderResult;
+use katana_provider::{ProviderFactory, ProviderResult};
 use katana_stage::trie::StateTrie;
 use katana_stage::{Stage, StageExecutionInput};
 use rstest::rstest;
 use starknet::macros::short_string;
 use starknet_types_core::hash::{Poseidon, StarkHash};
 
-/// Mock provider implementation for testing StateTrie stage.
+/// Mock ProviderFactory implementation for testing StateTrie stage.
 ///
 /// Provides configurable responses for headers, state updates, and trie operations.
 #[derive(Clone)]
 struct MockProvider {
-    /// Map of block number to header.
-    headers: Arc<Mutex<HashMap<BlockNumber, Header>>>,
-    /// Map of block number to state update.
-    state_updates: Arc<Mutex<HashMap<BlockNumber, StateUpdates>>>,
-    /// Track trie insert calls for verification.
-    trie_insert_calls: Arc<Mutex<Vec<BlockNumber>>>,
-    /// Whether to return an error on trie operations.
-    should_fail: Arc<Mutex<bool>>,
+    inner: MockInnerProvider,
+}
+
+impl std::fmt::Debug for MockProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MockProvider").finish_non_exhaustive()
+    }
 }
 
 impl MockProvider {
     fn new() -> Self {
         Self {
-            headers: Arc::new(Mutex::new(HashMap::new())),
-            state_updates: Arc::new(Mutex::new(HashMap::new())),
-            trie_insert_calls: Arc::new(Mutex::new(Vec::new())),
-            should_fail: Arc::new(Mutex::new(false)),
+            inner: MockInnerProvider {
+                headers: Arc::new(Mutex::new(HashMap::new())),
+                state_updates: Arc::new(Mutex::new(HashMap::new())),
+                trie_insert_calls: Arc::new(Mutex::new(Vec::new())),
+                should_fail: Arc::new(Mutex::new(false)),
+            },
         }
     }
 
     /// Configure a header for a specific block.
     fn with_header(self, block_number: BlockNumber, header: Header) -> Self {
-        self.headers.lock().unwrap().insert(block_number, header);
+        self.inner.headers.lock().unwrap().insert(block_number, header);
         self
     }
 
     /// Configure a state update for a specific block.
     fn with_state_update(self, block_number: BlockNumber, state_update: StateUpdates) -> Self {
-        self.state_updates.lock().unwrap().insert(block_number, state_update);
-        self
-    }
-
-    /// Configure the mock to fail on trie operations.
-    fn with_trie_error(self) -> Self {
-        *self.should_fail.lock().unwrap() = true;
+        self.inner.state_updates.lock().unwrap().insert(block_number, state_update);
         self
     }
 
     /// Get all block numbers that had trie inserts called.
     fn trie_insert_blocks(&self) -> Vec<BlockNumber> {
-        self.trie_insert_calls.lock().unwrap().clone()
+        self.inner.trie_insert_calls.lock().unwrap().clone()
     }
 }
 
-impl HeaderProvider for MockProvider {
+impl ProviderFactory for MockProvider {
+    type Provider = MockInnerProvider;
+    type ProviderMut = MockInnerProvider;
+
+    fn provider(&self) -> Self::Provider {
+        self.inner.clone()
+    }
+
+    fn provider_mut(&self) -> Self::ProviderMut {
+        self.inner.clone()
+    }
+}
+
+/// Mock inner provider implementation for testing StateTrie stage.
+///
+/// Provides configurable responses for headers, state updates, and trie operations.
+#[derive(Clone, Debug)]
+struct MockInnerProvider {
+    /// Map of block number to header.
+    headers: Arc<Mutex<HashMap<BlockNumber, Header>>>,
+    /// Map of block number to state update.
+    state_updates: Arc<Mutex<HashMap<BlockNumber, StateUpdates>>>,
+    /// Track trie insert calls for verification.
+    trie_insert_calls: Arc<Mutex<Vec<BlockNumber>>>,
+    /// Whether to return an error on trie operations.
+    should_fail: Arc<Mutex<bool>>,
+}
+
+impl HeaderProvider for MockInnerProvider {
     fn header(&self, id: BlockHashOrNumber) -> ProviderResult<Option<Header>> {
         let block_number = match id {
             BlockHashOrNumber::Num(num) => num,
@@ -80,7 +103,7 @@ impl HeaderProvider for MockProvider {
     }
 }
 
-impl StateUpdateProvider for MockProvider {
+impl StateUpdateProvider for MockInnerProvider {
     fn state_update(&self, block_id: BlockHashOrNumber) -> ProviderResult<Option<StateUpdates>> {
         let block_number = match block_id {
             BlockHashOrNumber::Num(num) => num,
@@ -109,7 +132,7 @@ impl StateUpdateProvider for MockProvider {
     }
 }
 
-impl TrieWriter for MockProvider {
+impl TrieWriter for MockInnerProvider {
     fn trie_insert_declared_classes(
         &self,
         block_number: BlockNumber,
@@ -139,6 +162,12 @@ impl TrieWriter for MockProvider {
     }
 }
 
+impl katana_provider::MutableProvider for MockInnerProvider {
+    fn commit(self) -> ProviderResult<()> {
+        Ok(())
+    }
+}
+
 /// Helper function to compute the expected state root from mock trie roots.
 fn compute_mock_state_root() -> Felt {
     let class_trie_root = Felt::from(0x1234u64);