feat: mock transport helpers

xJonathanLEI · Nov 26, 2023 · 320f82a · 320f82a
1 parent 5ca3b5f
commit 320f82a
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 25 deletions.
diff --git a/starknet-providers/src/jsonrpc/transports/mock.rs b/starknet-providers/src/jsonrpc/transports/mock.rs
@@ -16,9 +16,9 @@ use super::{HttpTransport, HttpTransportError};
 #[derive(Debug)]
 pub struct MockTransport {
     // Mock requests lookup
-    pub mock: HashMap<String, String>,
+    mocked_requests: HashMap<String, String>,
     // Mock method lookup if request lookup is None
-    pub method_mock: HashMap<String, String>,
+    mocked_methods: HashMap<String, String>,
     // Requests made
     pub requests_log: Arc<Mutex<Vec<(String, String)>>>,
     // HTTP fallback to help build mock requests
@@ -58,18 +58,27 @@ impl MockTransport {
     ///
     /// ```
     pub fn new(
-        mock: HashMap<String, String>,
-        method_mock: HashMap<String, String>,
         http_transport: Option<HttpTransport>,
         requests_log: Arc<Mutex<Vec<(String, String)>>>,
     ) -> Self {
         Self {
-            mock,
-            method_mock,
+            mocked_requests: HashMap::new(),
+            mocked_methods: HashMap::new(),
             requests_log,
             http_transport,
         }
     }
+
+    pub fn mock_request(&mut self, request_json: String, response_json: String) {
+        self.mocked_requests.insert(request_json, response_json);
+    }
+
+    pub fn mock_method(&mut self, method: JsonRpcMethod, response_json: String) {
+        let method_str = serde_json::to_string(&method)
+            .map_err(MockTransportError::Json)
+            .unwrap();
+        self.mocked_methods.insert(method_str, response_json);
+    }
 }
 
 #[cfg_attr(not(target_arch = "wasm32"), async_trait)]
@@ -100,17 +109,17 @@ impl JsonRpcTransport for MockTransport {
 
         let response_body;
 
-        if let Some(request_mock) = self.mock.get(&request_json) {
+        if let Some(request_mock) = self.mocked_requests.get(&request_json) {
             response_body = request_mock.clone();
-        } else if let Some(method_mock) = self.method_mock.get(&method_str) {
+        } else if let Some(method_mock) = self.mocked_methods.get(&method_str) {
             response_body = method_mock.clone();
         } else if let Some(http_transport) = &self.http_transport {
             response_body = http_transport
                 .send_request_raw(request_json.clone())
                 .await
                 .map_err(MockTransportError::Http)?;
             println!("\nUse this code to mock this request\n\n```rs");
-            println!("mock_transport.mock.insert(\n  {request_json:?}.into(), \n  {response_body:?}.into()\n);");
+            println!("mock_transport.mock_request(\n    {request_json:?}.into(),\n    {response_body:?}.into()\n);");
             // serde_json::to_string(&resp)?;
             println!("```\n");
         } else {

diff --git a/starknet-providers/tests/mock.rs b/starknet-providers/tests/mock.rs
@@ -5,7 +5,7 @@ use std::{
 
 use starknet_core::types::{BlockId, BlockTag, MaybePendingBlockWithTxHashes};
 use starknet_providers::{
-    jsonrpc::{HttpTransport, JsonRpcClient, MockTransport},
+    jsonrpc::{HttpTransport, JsonRpcClient, JsonRpcMethod, MockTransport},
     Provider,
 };
 use url::Url;
@@ -17,12 +17,7 @@ fn mock_transport_with_http() -> (Arc<Mutex<Vec<(String, String)>>>, MockTranspo
     let req_log = Arc::new(Mutex::new(vec![]));
     (
         req_log.clone(),
-        MockTransport::new(
-            HashMap::new(),
-            HashMap::new(),
-            Some(http_transport),
-            req_log,
-        ),
+        MockTransport::new(Some(http_transport), req_log),
     )
 }
 
@@ -43,22 +38,38 @@ async fn mock_transport_fallback() {
     };
 
     assert!(block.block_number > 0);
+}
+
+#[tokio::test]
+async fn mock_transport() {
+    let (_, mut mock_transport) = mock_transport_with_http();
+    // Block number 100000
+    mock_transport.mock_request(
+        "{\"id\":1,\"jsonrpc\":\"2.0\",\"method\":\"starknet_getBlockWithTxHashes\",\"params\":[\"latest\"]}".into(), 
+         "{\"jsonrpc\":\"2.0\",\"result\":{\"block_hash\":\"0x127edd99c58b5e7405c3fa24920abbf4c3fcfcd532a1c9f496afb917363c386\",\"block_number\":100000,\"new_root\":\"0x562df6c11a47b6711242d00318fec36c9f0f2613f7b711cd732857675b4f7f5\",\"parent_hash\":\"0x294f21cc482c8329b7e1f745cff69071685aec7955de7f5f9dae2be3cc27446\",\"sequencer_address\":\"0x1176a1bd84444c89232ec27754698e5d2e7e1a7f1539f12027f28b23ec9f3d8\",\"status\":\"ACCEPTED_ON_L2\",\"timestamp\":1701037710,\"transactions\":[\"0x1\"]},\"id\":1}".into()
+    );
 
-    // Test mock
+    let rpc_client = JsonRpcClient::new(mock_transport);
 
-    // Test method mock
+    let block = rpc_client
+        .get_block_with_tx_hashes(BlockId::Tag(BlockTag::Latest))
+        .await
+        .unwrap();
 
-    // Test fallback
+    let block = match block {
+        MaybePendingBlockWithTxHashes::Block(block) => block,
+        _ => panic!("unexpected block response type"),
+    };
 
-    // Test fallback helpful prints
+    assert!(block.block_number == 100000);
 }
 
 #[tokio::test]
-async fn mock_transport() {
+async fn mock_transport_method() {
     let (_, mut mock_transport) = mock_transport_with_http();
     // Block number 100000
-    mock_transport.mock.insert(
-        "{\"id\":1,\"jsonrpc\":\"2.0\",\"method\":\"starknet_getBlockWithTxHashes\",\"params\":[\"latest\"]}".into(), 
+    mock_transport.mock_method(
+        JsonRpcMethod::GetBlockWithTxHashes,
          "{\"jsonrpc\":\"2.0\",\"result\":{\"block_hash\":\"0x127edd99c58b5e7405c3fa24920abbf4c3fcfcd532a1c9f496afb917363c386\",\"block_number\":100000,\"new_root\":\"0x562df6c11a47b6711242d00318fec36c9f0f2613f7b711cd732857675b4f7f5\",\"parent_hash\":\"0x294f21cc482c8329b7e1f745cff69071685aec7955de7f5f9dae2be3cc27446\",\"sequencer_address\":\"0x1176a1bd84444c89232ec27754698e5d2e7e1a7f1539f12027f28b23ec9f3d8\",\"status\":\"ACCEPTED_ON_L2\",\"timestamp\":1701037710,\"transactions\":[\"0x1\"]},\"id\":1}".into()
       );
 
@@ -80,9 +91,9 @@ async fn mock_transport() {
 #[tokio::test]
 async fn mock_transport_log() {
     let (logs, mut mock_transport) = mock_transport_with_http();
-    mock_transport.mock.insert(
+    mock_transport.mock_request(
         "{\"id\":1,\"jsonrpc\":\"2.0\",\"method\":\"starknet_getBlockWithTxHashes\",\"params\":[\"latest\"]}".into(), 
-        "{\"jsonrpc\":\"2.0\",\"result\":{\"block_hash\":\"0x42fd8152ab51f0d5937ca83225035865c0dcdaea85ab84d38243ec5df23edac\",\"block_number\":100000,\"new_root\":\"0x372c133dace5d2842e3791741b6c05af840f249b52febb18f483d1eb38aaf8a\",\"parent_hash\":\"0x7f6df65f94584de3ff9807c67822197692cc8895aa1de5340af0072ac2ccfb5\",\"sequencer_address\":\"0x1176a1bd84444c89232ec27754698e5d2e7e1a7f1539f12027f28b23ec9f3d8\",\"status\":\"ACCEPTED_ON_L2\",\"timestamp\":1701033987,\"transactions\":[\"0x1\",\"0x2\",\"0x3\"]},\"id\":1}".into()
+        "{\"jsonrpc\":\"2.0\",\"result\":{\"block_hash\":\"0x42fd8152ab51f0d5937ca83225035865c0dcdaea85ab84d38243ec5df23edac\",\"block_number\":100000,\"new_root\":\"0x372c133dace5d2842e3791741b6c05af840f249b52febb18f483d1eb38aaf8a\",\"parent_hash\":\"0x7f6df65f94584de3ff9807c67822197692cc8895aa1de5340af0072ac2ccfb5\",\"sequencer_address\":\"0x1176a1bd84444c89232ec27754698e5d2e7e1a7f1539f12027f28b23ec9f3d8\",\"status\":\"ACCEPTED_ON_L2\",\"timestamp\":1701033987,\"transactions\":[\"0x1\"]},\"id\":1}".into()
       );
 
     let rpc_client = JsonRpcClient::new(mock_transport);