From 219d00d3ecef914eb44ae1090525e8d704377889 Mon Sep 17 00:00:00 2001
From: Alex Ostrovski <aov@matterlabs.dev>
Date: Mon, 6 Nov 2023 14:38:55 +0200
Subject: [PATCH] Update `SyncBlocks`-related tests

---
 node/actors/executor/src/tests.rs               | 16 +++-------------
 node/actors/sync_blocks/src/peers/tests.rs      | 11 ++++-------
 node/actors/sync_blocks/src/tests/end_to_end.rs | 15 ++++++---------
 node/libs/crypto/src/bn254/tests.rs             |  6 ++----
 4 files changed, 15 insertions(+), 33 deletions(-)

diff --git a/node/actors/executor/src/tests.rs b/node/actors/executor/src/tests.rs
index 64796997..606818a8 100644
--- a/node/actors/executor/src/tests.rs
+++ b/node/actors/executor/src/tests.rs
@@ -6,16 +6,6 @@ use concurrency::sync;
 use roles::validator::{BlockNumber, Payload};
 use storage::{BlockStore, InMemoryStorage, StorageError};
 
-async fn run_executor(ctx: &ctx::Ctx, executor: Executor<InMemoryStorage>) -> anyhow::Result<()> {
-    executor.run(ctx).await.or_else(|err| {
-        if err.root_cause().is::<ctx::Canceled>() {
-            Ok(()) // Test has successfully finished
-        } else {
-            Err(err)
-        }
-    })
-}
-
 async fn store_final_blocks(
     ctx: &ctx::Ctx,
     mut blocks_receiver: channel::UnboundedReceiver<FinalBlock>,
@@ -69,7 +59,7 @@ async fn executing_single_validator() {
     let (executor, mut blocks_receiver) = validator.into_executor(storage);
 
     scope::run!(ctx, |ctx, s| async {
-        s.spawn_bg(run_executor(ctx, executor));
+        s.spawn_bg(executor.run(ctx));
 
         let mut expected_block_number = BlockNumber(1);
         while expected_block_number < BlockNumber(5) {
@@ -109,8 +99,8 @@ async fn executing_validator_and_external_node() {
     .unwrap();
 
     scope::run!(ctx, |ctx, s| async {
-        s.spawn_bg(run_executor(ctx, validator));
-        s.spawn_bg(run_executor(ctx, external_node));
+        s.spawn_bg(validator.run(ctx));
+        s.spawn_bg(external_node.run(ctx));
         s.spawn_bg(store_final_blocks(ctx, blocks_receiver, validator_storage));
 
         for _ in 0..5 {
diff --git a/node/actors/sync_blocks/src/peers/tests.rs b/node/actors/sync_blocks/src/peers/tests.rs
index 9dcf01ad..6eea24fa 100644
--- a/node/actors/sync_blocks/src/peers/tests.rs
+++ b/node/actors/sync_blocks/src/peers/tests.rs
@@ -6,7 +6,7 @@ use concurrency::time;
 use rand::{rngs::StdRng, seq::IteratorRandom, Rng};
 use roles::validator;
 use std::{collections::HashSet, fmt};
-use storage::{BlockStore, InMemoryStorage};
+use storage::{BlockStore, InMemoryStorage, StorageError};
 use test_casing::{test_casing, Product};
 
 const TEST_TIMEOUT: time::Duration = time::Duration::seconds(5);
@@ -115,12 +115,9 @@ async fn test_peer_states<T: Test>(test: T) {
 
     scope::run!(ctx, |ctx, s| async {
         s.spawn_bg(async {
-            peer_states.run(ctx).await.or_else(|err| {
-                if err.root_cause().is::<ctx::Canceled>() {
-                    Ok(()) // Swallow cancellation errors after the test is finished
-                } else {
-                    Err(err)
-                }
+            peer_states.run(ctx).await.or_else(|err| match err {
+                StorageError::Canceled(_) => Ok(()), // Swallow cancellation errors after the test is finished
+                StorageError::Database(err) => Err(err),
             })
         });
         test.test(ctx, test_handles).await
diff --git a/node/actors/sync_blocks/src/tests/end_to_end.rs b/node/actors/sync_blocks/src/tests/end_to_end.rs
index 7bb78b1f..1d9256bf 100644
--- a/node/actors/sync_blocks/src/tests/end_to_end.rs
+++ b/node/actors/sync_blocks/src/tests/end_to_end.rs
@@ -115,7 +115,7 @@ impl Node {
             sync_blocks_config,
         )
         .await
-        .expect("Failed");
+        .expect("Failed initializing `sync_blocks` actor");
 
         let sync_states_subscriber = sync_blocks.subscribe_to_state_updates();
         self.network
@@ -158,8 +158,10 @@ impl Node {
 
             self.switch_off_receiver
                 .recv_or_disconnected(ctx)
-                .await?
+                .await
                 .ok();
+            // ^ Unlike with `switch_on_receiver`, the context may get canceled before the receiver
+            // is dropped, so we swallow both cancellation and disconnect errors here.
             tracing::trace!("Node stopped");
             Ok(())
         })
@@ -240,14 +242,9 @@ async fn test_sync_blocks<T: GossipNetworkTest>(test: T) {
             s.spawn_bg(async {
                 let test_validators = test_validators;
                 let key = node.key();
-                let err = node.run(ctx, &test_validators).await.unwrap_err();
-
+                node.run(ctx, &test_validators).await?;
                 tracing::trace!(?key, "Node task completed");
-                if err.root_cause().is::<ctx::Canceled>() {
-                    Ok(()) // Test has successfully completed
-                } else {
-                    Err(err)
-                }
+                Ok(())
             });
         }
 
diff --git a/node/libs/crypto/src/bn254/tests.rs b/node/libs/crypto/src/bn254/tests.rs
index 933b914b..d07afbfa 100644
--- a/node/libs/crypto/src/bn254/tests.rs
+++ b/node/libs/crypto/src/bn254/tests.rs
@@ -1,8 +1,6 @@
-use std::iter::repeat_with;
-
-use rand::{rngs::StdRng, Rng, SeedableRng};
-
 use crate::bn254::{AggregateSignature, PublicKey, SecretKey, Signature};
+use rand::{rngs::StdRng, Rng, SeedableRng};
+use std::iter::repeat_with;
 
 #[test]
 fn signature_smoke() {