fix: unit-tests fixed on main (matter-labs#193)

src/cache.rs

@@ -50,6 +50,7 @@ pub(crate) struct Cache {
     block_raw_transactions: FxHashMap<u64, Vec<RawTransaction>>,
     transactions: FxHashMap<H256, Transaction>,
     bridge_addresses: Option<BridgeAddresses>,
+    confirmed_tokens: FxHashMap<(u32, u8), Vec<zksync_web3_decl::types::Token>>,
 }
 
 impl Cache {
@@ -158,6 +159,32 @@ impl Cache {
         self.block_raw_transactions.get(number)
     }
 
+    /// Returns the cached confirmed tokens.
+    pub(crate) fn get_confirmed_tokens(
+        &self,
+        from: u32,
+        limit: u8,
+    ) -> Option<&Vec<zksync_web3_decl::types::Token>> {
+        if matches!(self.config, CacheConfig::None) {
+            return None;
+        }
+        self.confirmed_tokens.get(&(from, limit))
+    }
+
+    /// Cache confirmed tokens
+    pub(crate) fn set_confirmed_tokens(
+        &mut self,
+        from: u32,
+        limit: u8,
+        confirmed_tokens: Vec<zksync_web3_decl::types::Token>,
+    ) {
+        if matches!(self.config, CacheConfig::None) {
+            return;
+        }
+        self.confirmed_tokens
+            .insert((from, limit), confirmed_tokens);
+    }
+
     /// Cache the raw transactions for the provided block number.
     pub(crate) fn insert_block_raw_transactions(
         &mut self,

src/fork.rs

@@ -260,6 +260,13 @@ pub trait ForkSource {
 
     /// Returns addresses of the default bridge contracts.
     fn get_bridge_contracts(&self) -> eyre::Result<BridgeAddresses>;
+
+    /// Returns confirmed tokens
+    fn get_confirmed_tokens(
+        &self,
+        from: u32,
+        limit: u8,
+    ) -> eyre::Result<Vec<zksync_web3_decl::types::Token>>;
 }
 
 /// Holds the information about the original chain.

src/http_fork_source.rs

@@ -1,19 +1,19 @@
 use std::sync::RwLock;
 
+use crate::{
+    cache::{Cache, CacheConfig},
+    fork::{block_on, ForkSource},
+};
 use eyre::Context;
 use zksync_basic_types::{H256, U256};
 use zksync_types::api::{BridgeAddresses, Transaction};
+use zksync_web3_decl::types::Token;
 use zksync_web3_decl::{
     jsonrpsee::http_client::{HttpClient, HttpClientBuilder},
     namespaces::{EthNamespaceClient, ZksNamespaceClient},
     types::Index,
 };
 
-use crate::{
-    cache::{Cache, CacheConfig},
-    fork::{block_on, ForkSource},
-};
-
 #[derive(Debug)]
 /// Fork source that gets the data via HTTP requests.
 pub struct HttpForkSource {
@@ -306,6 +306,37 @@ impl ForkSource for HttpForkSource {
             })
             .wrap_err("fork http client failed")
     }
+
+    /// Returns known token addresses
+    fn get_confirmed_tokens(&self, from: u32, limit: u8) -> eyre::Result<Vec<Token>> {
+        if let Some(confirmed_tokens) = self
+            .cache
+            .read()
+            .ok()
+            .and_then(|guard| guard.get_confirmed_tokens(from, limit).cloned())
+        {
+            tracing::debug!("using cached confirmed_tokens");
+            return Ok(confirmed_tokens);
+        };
+
+        let client = self.create_client();
+        block_on(async move { client.get_confirmed_tokens(from, limit).await })
+            .map(|confirmed_tokens| {
+                self.cache
+                    .write()
+                    .map(|mut guard| {
+                        guard.set_confirmed_tokens(from, limit, confirmed_tokens.clone())
+                    })
+                    .unwrap_or_else(|err| {
+                        tracing::warn!(
+                            "failed writing to cache for 'set_confirmed_tokens': {:?}",
+                            err
+                        )
+                    });
+                confirmed_tokens
+            })
+            .wrap_err("fork http client failed")
+    }
 }
 
 #[cfg(test)]

src/zks.rs

@@ -1,8 +1,11 @@
-use std::sync::{Arc, RwLock};
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
 
 use bigdecimal::BigDecimal;
 use futures::FutureExt;
-use zksync_basic_types::{Address, L1BatchNumber, MiniblockNumber, U256};
+use zksync_basic_types::{AccountTreeId, Address, L1BatchNumber, MiniblockNumber, U256};
 use zksync_core::api_server::web3::backend_jsonrpc::{
     error::{internal_error, into_jsrpc_error},
     namespaces::zks::ZksNamespaceT,
@@ -14,8 +17,10 @@ use zksync_types::{
         TransactionDetails, TransactionStatus, TransactionVariant,
     },
     fee::Fee,
+    utils::storage_key_for_standard_token_balance,
     ExecuteTransactionCommon, ProtocolVersionId, Transaction,
 };
+use zksync_utils::h256_to_u256;
 use zksync_web3_decl::{
     error::Web3Error,
     types::{Filter, Log},
@@ -200,10 +205,29 @@ impl<S: Send + Sync + 'static + ForkSource + std::fmt::Debug> ZksNamespaceT
 
     fn get_confirmed_tokens(
         &self,
-        _from: u32,
-        _limit: u8,
+        from: u32,
+        limit: u8,
     ) -> jsonrpc_core::BoxFuture<jsonrpc_core::Result<Vec<zksync_web3_decl::types::Token>>> {
-        not_implemented("zks_getConfirmedTokens")
+        let inner = self.node.clone();
+        Box::pin(async move {
+            let reader = inner
+                .read()
+                .map_err(|_| into_jsrpc_error(Web3Error::InternalError))?;
+
+            let fork_storage_read = reader
+                .fork_storage
+                .inner
+                .read()
+                .expect("failed reading fork storage");
+
+            match fork_storage_read.fork.as_ref() {
+                Some(fork) => Ok(fork
+                    .fork_source
+                    .get_confirmed_tokens(from, limit)
+                    .map_err(|_e| into_jsrpc_error(Web3Error::InternalError))?),
+                None => Ok(vec![]),
+            }
+        })
     }
 
     fn get_token_price(
@@ -241,13 +265,49 @@ impl<S: Send + Sync + 'static + ForkSource + std::fmt::Debug> ZksNamespaceT
         }
     }
 
+    /// Get all known balances for a given account.
+    ///
+    /// # Arguments
+    ///
+    /// * `address` - The user address with balances to check.
+    ///
+    /// # Returns
+    ///
+    /// A `BoxFuture` containing a `Result` with a (Token, Balance) map where account has non-zero value.
     fn get_all_account_balances(
         &self,
-        _address: zksync_basic_types::Address,
+        address: zksync_basic_types::Address,
     ) -> jsonrpc_core::BoxFuture<
         jsonrpc_core::Result<std::collections::HashMap<zksync_basic_types::Address, U256>>,
     > {
-        not_implemented("zks_getAllAccountBalances")
+        let inner = self.node.clone();
+        Box::pin({
+            let address = address.clone();
+            self.get_confirmed_tokens(0, 100)
+                .then(move |tokens| async move {
+                    let tokens =
+                        tokens.map_err(|_err| into_jsrpc_error(Web3Error::InternalError))?;
+
+                    let mut writer = inner
+                        .write()
+                        .map_err(|_err| into_jsrpc_error(Web3Error::InternalError))?;
+
+                    let mut balances = HashMap::new();
+                    for token in tokens {
+                        let balance_key = storage_key_for_standard_token_balance(
+                            AccountTreeId::new(token.l2_address),
+                            &address,
+                        );
+
+                        let balance = writer.fork_storage.read_value(&balance_key);
+                        if !balance.is_zero() {
+                            balances.insert(token.l2_address, h256_to_u256(balance));
+                        }
+                    }
+
+                    Ok(balances)
+                })
+        })
     }
 
     fn get_l2_to_l1_msg_proof(
@@ -1014,12 +1074,8 @@ mod tests {
 
         let node = InMemoryNode::<HttpForkSource>::new(
             Some(ForkDetails::from_network(&mock_server.url(), None, CacheConfig::None).await),
-            crate::node::ShowCalls::None,
-            ShowStorageLogs::None,
-            ShowVMDetails::None,
-            ShowGasDetails::None,
-            false,
-            &system_contracts::Options::BuiltIn,
+            None,
+            Default::default(),
         );
 
         let namespace = ZkMockNamespaceImpl::new(node.get_inner());

test_endpoints.http

@@ -729,6 +729,28 @@ content-type: application/json
 POST http://localhost:8011
 content-type: application/json
 
+{
+    "jsonrpc": "2.0",
+    "id": "1",
+    "method": "zks_getConfirmedTokens",
+    "params": [0, 100]
+}
+
+###
+POST http://localhost:8011
+content-type: application/json
+
+{
+    "jsonrpc": "2.0",
+    "id": "1",
+    "method": "zks_getAllAccountBalances",
+    "params": ["0x364d6D0333432C3Ac016Ca832fb8594A8cE43Ca6"]
+}
+
+###
+POST http://localhost:8011
+content-type: application/json
+
 {
     "jsonrpc": "2.0",
     "id": "1",