add(scan): Test the RegisterKeys scan service call (#8281)

* Impl generating continuous deserialized blocks

* Make `sapling_efvk_hrp` `pub`

* Don't wait for Sapling activation height in tests

* Set the sleep interval for scan service to 10 secs

* Simplify `sapling_key_to_scan_block_keys`

* Enable mocking Sapling scanning keys for Testnet

* Test the `RegisterKeys` scan service call

* Enable `shielded-scan` for `zebra-chain`

* Use an ephemeral database so results don't persist

* Don't generate blocks when mocking the state

* Improve error messages

* Simplify seeding mocked Sapling viewing keys

* Apply suggestions from code review

Co-authored-by: Arya <aryasolhi@gmail.com>

* Use a manual iterator over `Network`

---------

Co-authored-by: Arya <aryasolhi@gmail.com>

zebra-chain/src/parameters/network.rs

@@ -85,6 +85,12 @@ impl fmt::Display for Network {
 }
 
 impl Network {
+    /// Returns an iterator over [`Network`] variants.
+    pub fn iter() -> impl Iterator<Item = Self> {
+        // TODO: Use default values of `Testnet` variant when adding fields for #7845.
+        [Self::Mainnet, Self::Testnet].into_iter()
+    }
+
     /// Get the default port associated to this network.
     pub fn default_port(&self) -> u16 {
         match self {

zebra-chain/src/primitives/viewing_key/sapling.rs

@@ -74,7 +74,7 @@ impl SaplingViewingKey {
 impl Network {
     /// Returns the human-readable prefix for an Zcash Sapling extended full viewing key
     /// for this network.
-    fn sapling_efvk_hrp(&self) -> &'static str {
+    pub fn sapling_efvk_hrp(&self) -> &'static str {
         if self.is_a_test_network() {
             // Assume custom testnets have the same HRP
             //

zebra-scan/Cargo.toml

@@ -54,7 +54,7 @@ futures = "0.3.30"
 zcash_client_backend = "0.10.0-rc.1"
 zcash_primitives = "0.13.0-rc.1"
 
-zebra-chain = { path = "../zebra-chain", version = "1.0.0-beta.34" }
+zebra-chain = { path = "../zebra-chain", version = "1.0.0-beta.34", features = ["shielded-scan"] }
 zebra-state = { path = "../zebra-state", version = "1.0.0-beta.34", features = ["shielded-scan"] }
 zebra-node-services = { path = "../zebra-node-services", version = "1.0.0-beta.34", features = ["shielded-scan"] }
 zebra-grpc = { path = "../zebra-grpc", version = "0.1.0-alpha.1" }
@@ -75,7 +75,6 @@ zcash_note_encryption = { version = "0.4.0", optional = true }
 zebra-test = { path = "../zebra-test", version = "1.0.0-beta.34", optional = true }
 
 [dev-dependencies]
-
 insta = { version = "1.33.0", features = ["ron", "redactions"] }
 tokio = { version = "1.36.0", features = ["test-util"] }
 

zebra-scan/src/service/scan_task/scan.rs

@@ -21,7 +21,6 @@ use zcash_client_backend::{
     scanning::{ScanError, ScanningKey},
 };
 use zcash_primitives::{
-    constants::*,
     sapling::SaplingIvk,
     zip32::{AccountId, DiversifiableFullViewingKey, Scope},
 };
@@ -61,8 +60,9 @@ const INITIAL_WAIT: Duration = Duration::from_secs(15);
 
 /// The amount of time between checking for new blocks and starting new scans.
 ///
-/// This is just under half the target block interval.
-pub const CHECK_INTERVAL: Duration = Duration::from_secs(30);
+/// TODO: The current value is set to 10 so that tests don't sleep for too long and finish faster.
+///       Set it to 30 after #8250 gets addressed or remove this const completely in the refactor.
+pub const CHECK_INTERVAL: Duration = Duration::from_secs(10);
 
 /// We log an info log with progress after this many blocks.
 const INFO_LOG_INTERVAL: u32 = 10_000;
@@ -81,6 +81,7 @@ pub async fn start(
     info!(?network, "starting scan task");
 
     // Do not scan and notify if we are below sapling activation height.
+    #[cfg(not(test))]
     wait_for_height(
         sapling_activation_height,
         "Sapling activation",
@@ -405,19 +406,11 @@ pub fn scan_block<K: ScanningKey>(
 //       performance: stop returning both the dfvk and ivk for the same key
 // TODO: use `ViewingKey::parse` from zebra-chain instead
 pub fn sapling_key_to_scan_block_keys(
-    sapling_key: &SaplingScanningKey,
+    key: &SaplingScanningKey,
     network: Network,
 ) -> Result<(Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>), Report> {
-    let hrp = if network.is_a_test_network() {
-        // Assume custom testnets have the same HRP
-        //
-        // TODO: add the regtest HRP here
-        testnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY
-    } else {
-        mainnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY
-    };
-
-    let efvk = decode_extended_full_viewing_key(hrp, sapling_key).map_err(|e| eyre!(e))?;
+    let efvk =
+        decode_extended_full_viewing_key(network.sapling_efvk_hrp(), key).map_err(|e| eyre!(e))?;
 
     // Just return all the keys for now, so we can be sure our code supports them.
     let dfvk = efvk.to_diversifiable_full_viewing_key();

zebra-scan/src/service/scan_task/tests/vectors.rs

@@ -18,7 +18,8 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Send some keys to be registered
     let num_keys = 10;
-    let sapling_keys = mock_sapling_scanning_keys(num_keys.try_into().expect("should fit in u8"));
+    let sapling_keys =
+        mock_sapling_scanning_keys(num_keys.try_into().expect("should fit in u8"), network);
     let sapling_keys_with_birth_heights: Vec<(String, Option<u32>)> =
         sapling_keys.into_iter().zip((0..).map(Some)).collect();
     mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
@@ -60,7 +61,7 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Check that keys can't be overridden.
 
-    let sapling_keys = mock_sapling_scanning_keys(20);
+    let sapling_keys = mock_sapling_scanning_keys(20, network);
     let sapling_keys_with_birth_heights: Vec<(String, Option<u32>)> = sapling_keys
         .clone()
         .into_iter()
@@ -87,7 +88,7 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Check that it removes keys correctly
 
-    let sapling_keys = mock_sapling_scanning_keys(30);
+    let sapling_keys = mock_sapling_scanning_keys(30, network);
     let done_rx = mock_scan_task.remove_keys(&sapling_keys)?;
 
     let (new_keys, _new_results_senders) =

zebra-scan/src/service/tests.rs

@@ -1,7 +1,7 @@
 //! Tests for ScanService.
 
 use tokio::sync::mpsc::error::TryRecvError;
-use tower::{Service, ServiceExt};
+use tower::{Service, ServiceBuilder, ServiceExt};
 
 use color_eyre::{eyre::eyre, Result};
 
@@ -12,7 +12,8 @@ use zebra_state::TransactionIndex;
 use crate::{
     service::{scan_task::ScanTaskCommand, ScanService},
     storage::db::tests::{fake_sapling_results, new_test_storage},
-    tests::ZECPAGES_SAPLING_VIEWING_KEY,
+    tests::{mock_sapling_scanning_keys, ZECPAGES_SAPLING_VIEWING_KEY},
+    Config,
 };
 
 /// Tests that keys are deleted correctly
@@ -254,3 +255,78 @@ pub async fn scan_service_get_results_for_key_correctly() -> Result<()> {
 
     Ok(())
 }
+
+/// Tests that the scan service registers keys correctly.
+#[tokio::test]
+pub async fn scan_service_registers_keys_correctly() -> Result<()> {
+    for network in Network::iter() {
+        scan_service_registers_keys_correctly_for(network).await?;
+    }
+
+    Ok(())
+}
+
+async fn scan_service_registers_keys_correctly_for(network: Network) -> Result<()> {
+    // Mock the state.
+    let (state, _, _, chain_tip_change) = zebra_state::populated_state(vec![], network).await;
+
+    // Instantiate the scan service.
+    let mut scan_service = ServiceBuilder::new()
+        .buffer(2)
+        .service(ScanService::new(&Config::ephemeral(), network, state, chain_tip_change).await);
+
+    // Mock three Sapling keys.
+    let mocked_keys = mock_sapling_scanning_keys(3, network);
+
+    // Add birth heights to the mocked keys.
+    let keys_to_register: Vec<_> = mocked_keys
+        .clone()
+        .into_iter()
+        .zip((0u32..).map(Some))
+        .collect();
+
+    // Register the first key.
+    match scan_service
+        .ready()
+        .await
+        .map_err(|err| eyre!(err))?
+        .call(Request::RegisterKeys(keys_to_register[..1].to_vec()))
+        .await
+        .map_err(|err| eyre!(err))?
+    {
+        Response::RegisteredKeys(registered_keys) => {
+            // The key should be registered.
+            assert_eq!(
+                registered_keys,
+                mocked_keys[..1],
+                "response should match newly registered key"
+            );
+        }
+
+        _ => panic!("scan service should have responded with the `RegisteredKeys` response"),
+    }
+
+    // Try registering all three keys.
+    match scan_service
+        .ready()
+        .await
+        .map_err(|err| eyre!(err))?
+        .call(Request::RegisterKeys(keys_to_register))
+        .await
+        .map_err(|err| eyre!(err))?
+    {
+        Response::RegisteredKeys(registered_keys) => {
+            // Only the last two keys should be registered in this service call since the first one
+            // was registered in the previous call.
+            assert_eq!(
+                registered_keys,
+                mocked_keys[1..3],
+                "response should match newly registered keys"
+            );
+        }
+
+        _ => panic!("scan service should have responded with the `RegisteredKeys` response"),
+    }
+
+    Ok(())
+}

zebra-scan/src/tests.rs

@@ -37,6 +37,7 @@ use zebra_chain::{
     amount::{Amount, NegativeAllowed},
     block::{self, merkle, Block, Header, Height},
     fmt::HexDebug,
+    parameters::Network,
     primitives::{redjubjub, Groth16Proof},
     sapling::{self, PerSpendAnchor, Spend, TransferData},
     serialization::AtLeastOne,
@@ -55,16 +56,16 @@ pub const ZECPAGES_SAPLING_VIEWING_KEY: &str = "zxviews1q0duytgcqqqqpqre26wkl45g
 /// A fake viewing key in an incorrect format.
 pub const FAKE_SAPLING_VIEWING_KEY: &str = "zxviewsfake";
 
-/// Generates `num_keys` of [`SaplingScanningKey`]s for tests.
+/// Generates `num_keys` of [`SaplingScanningKey`]s for tests for the given [`Network`].
 ///
 /// The keys are seeded only from their index in the returned `Vec`, so repeated calls return same
 /// keys at a particular index.
-pub fn mock_sapling_scanning_keys(num_keys: u8) -> Vec<SaplingScanningKey> {
+pub fn mock_sapling_scanning_keys(num_keys: u8, network: Network) -> Vec<SaplingScanningKey> {
     let mut keys: Vec<SaplingScanningKey> = vec![];
 
     for seed in 0..num_keys {
         keys.push(encode_extended_full_viewing_key(
-            "zxviews",
+            network.sapling_efvk_hrp(),
             &mock_sapling_efvk(&[seed]),
         ));
     }