Use mockall mocks for block remover tests, instead of handwritten (My…

…stenLabs#5065)

narwhal/primary/src/tests/block_remover_tests.rs

@@ -1,10 +1,7 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
-use crate::{block_remover::BlockRemover, common::create_db_stores, PrimaryWorkerMessage};
-
-use anemo::async_trait;
+use crate::{block_remover::BlockRemover, common::create_db_stores};
 use anemo::PeerId;
-use anyhow::Result;
 use config::{Committee, WorkerId};
 use consensus::{dag::Dag, metrics::ConsensusMetrics};
 use crypto::traits::KeyPair;
@@ -14,11 +11,10 @@ use network::P2pNetwork;
 use prometheus::Registry;
 use std::{borrow::Borrow, collections::HashMap, sync::Arc, time::Duration};
 use test_utils::CommitteeFixture;
-use test_utils::WorkerFixture;
 use tokio::time::timeout;
 use types::{
-    BatchDigest, Certificate, PrimaryToWorker, PrimaryToWorkerServer, RequestBatchRequest,
-    RequestBatchResponse, WorkerDeleteBatchesMessage, WorkerSynchronizeMessage,
+    BatchDigest, Certificate, MockPrimaryToWorker, PrimaryToWorkerServer,
+    WorkerDeleteBatchesMessage,
 };
 
 #[tokio::test]
@@ -112,7 +108,19 @@ async fn test_successful_blocks_delete() {
         let worker = primary.worker(worker_id);
         let address = &worker.info().worker_address;
 
-        worker_networks.push(worker_listener(worker, batch_digests, Ok(())));
+        let mut mock_server = MockPrimaryToWorker::new();
+        mock_server
+            .expect_delete_batches()
+            .withf(move |request| {
+                request.body()
+                    == &WorkerDeleteBatchesMessage {
+                        digests: batch_digests.clone(),
+                    }
+            })
+            .returning(|_| Ok(anemo::Response::new(())));
+        let routes = anemo::Router::new().add_rpc_service(PrimaryToWorkerServer::new(mock_server));
+        worker_networks.push(worker.new_network(routes));
+        // worker_networks.push(worker_listener(worker, batch_digests, Ok(())));
 
         let address = network::multiaddr_to_address(address).unwrap();
         let peer_id = PeerId(worker.keypair().public().0.to_bytes());
@@ -261,15 +269,24 @@ async fn test_failed_blocks_delete() {
         let worker = primary.worker(worker_id);
         let address = &worker.info().worker_address;
 
-        worker_networks.push(worker_listener(
-            worker,
-            batch_digests,
-            if worker_id == 0 {
-                Err(anyhow::anyhow!("failed"))
-            } else {
-                Ok(())
-            },
-        ));
+        let mut mock_server = MockPrimaryToWorker::new();
+        mock_server
+            .expect_delete_batches()
+            .withf(move |request| {
+                request.body()
+                    == &WorkerDeleteBatchesMessage {
+                        digests: batch_digests.clone(),
+                    }
+            })
+            .returning(move |_| {
+                if worker_id == 0 {
+                    Err(anemo::rpc::Status::internal("failed"))
+                } else {
+                    Ok(anemo::Response::new(()))
+                }
+            });
+        let routes = anemo::Router::new().add_rpc_service(PrimaryToWorkerServer::new(mock_server));
+        worker_networks.push(worker.new_network(routes));
 
         let address = network::multiaddr_to_address(address).unwrap();
         let peer_id = PeerId(worker.keypair().public().0.to_bytes());
@@ -307,61 +324,6 @@ async fn test_failed_blocks_delete() {
     assert_eq!(total_deleted, 0);
 }
 
-#[must_use]
-pub fn worker_listener(
-    worker: &WorkerFixture,
-    expected_batch_ids: Vec<BatchDigest>,
-    response: Result<()>,
-) -> anemo::Network {
-    struct MockPrimaryToWorker {
-        expected_request: Vec<BatchDigest>,
-        response: Result<()>,
-    }
-    #[async_trait]
-    impl PrimaryToWorker for MockPrimaryToWorker {
-        async fn send_message(
-            &self,
-            _request: anemo::Request<PrimaryWorkerMessage>,
-        ) -> Result<anemo::Response<()>, anemo::rpc::Status> {
-            unimplemented!()
-        }
-        async fn synchronize(
-            &self,
-            _request: anemo::Request<WorkerSynchronizeMessage>,
-        ) -> Result<anemo::Response<()>, anemo::rpc::Status> {
-            unimplemented!()
-        }
-        async fn request_batch(
-            &self,
-            _request: anemo::Request<RequestBatchRequest>,
-        ) -> Result<anemo::Response<RequestBatchResponse>, anemo::rpc::Status> {
-            unimplemented!()
-        }
-        async fn delete_batches(
-            &self,
-            request: anemo::Request<WorkerDeleteBatchesMessage>,
-        ) -> Result<anemo::Response<()>, anemo::rpc::Status> {
-            assert_eq!(
-                request.into_body(),
-                WorkerDeleteBatchesMessage {
-                    digests: self.expected_request.clone()
-                }
-            );
-            match &self.response {
-                Ok(_) => Ok(anemo::Response::new(())),
-                Err(e) => Err(anemo::rpc::Status::unknown(format!("{e:?}"))),
-            }
-        }
-    }
-
-    let routes =
-        anemo::Router::new().add_rpc_service(PrimaryToWorkerServer::new(MockPrimaryToWorker {
-            expected_request: expected_batch_ids,
-            response,
-        }));
-    worker.new_network(routes)
-}
-
 async fn populate_genesis<K: Borrow<Dag>>(dag: &K, committee: &Committee) {
     assert!(join_all(
         Certificate::genesis(committee)

narwhal/types/Cargo.toml

@@ -15,6 +15,7 @@ dashmap = "5.4.0"
 derive_builder = "0.11.2"
 futures = "0.3.24"
 indexmap = { version = "1.9.1", features = ["serde"] }
+mockall = "0.11.2"
 mysten-util-mem = { git = "https://github.com/MystenLabs/mysten-infra" }
 prometheus = "0.13.2"
 proptest = "1.0.0"

narwhal/types/build.rs

@@ -42,9 +42,13 @@ fn main() -> Result<()> {
 }
 
 fn build_anemo_services(out_dir: &Path) {
+    let mut automock_attribute = anemo_build::Attributes::default();
+    automock_attribute.push_trait(".", r#"#[mockall::automock]"#);
+
     let primary_to_primary = anemo_build::manual::Service::builder()
         .name("PrimaryToPrimary")
         .package("narwhal")
+        .attributes(automock_attribute.clone())
         .method(
             anemo_build::manual::Method::builder()
                 .name("send_message")
@@ -59,6 +63,7 @@ fn build_anemo_services(out_dir: &Path) {
     let primary_to_worker = anemo_build::manual::Service::builder()
         .name("PrimaryToWorker")
         .package("narwhal")
+        .attributes(automock_attribute.clone())
         .method(
             anemo_build::manual::Method::builder()
                 .name("send_message")
@@ -100,6 +105,7 @@ fn build_anemo_services(out_dir: &Path) {
     let worker_to_primary = anemo_build::manual::Service::builder()
         .name("WorkerToPrimary")
         .package("narwhal")
+        .attributes(automock_attribute.clone())
         .method(
             anemo_build::manual::Method::builder()
                 .name("send_message")
@@ -123,6 +129,7 @@ fn build_anemo_services(out_dir: &Path) {
     let worker_to_worker = anemo_build::manual::Service::builder()
         .name("WorkerToWorker")
         .package("narwhal")
+        .attributes(automock_attribute)
         .method(
             anemo_build::manual::Method::builder()
                 .name("send_message")

narwhal/types/src/proto.rs

@@ -22,19 +22,19 @@ pub use narwhal::{
     configuration_client::ConfigurationClient,
     configuration_server::{Configuration, ConfigurationServer},
     primary_to_primary_client::PrimaryToPrimaryClient,
-    primary_to_primary_server::{PrimaryToPrimary, PrimaryToPrimaryServer},
+    primary_to_primary_server::{MockPrimaryToPrimary, PrimaryToPrimary, PrimaryToPrimaryServer},
     primary_to_worker_client::PrimaryToWorkerClient,
-    primary_to_worker_server::{PrimaryToWorker, PrimaryToWorkerServer},
+    primary_to_worker_server::{MockPrimaryToWorker, PrimaryToWorker, PrimaryToWorkerServer},
     proposer_client::ProposerClient,
     proposer_server::{Proposer, ProposerServer},
     transactions_client::TransactionsClient,
     transactions_server::{Transactions, TransactionsServer},
     validator_client::ValidatorClient,
     validator_server::{Validator, ValidatorServer},
     worker_to_primary_client::WorkerToPrimaryClient,
-    worker_to_primary_server::{WorkerToPrimary, WorkerToPrimaryServer},
+    worker_to_primary_server::{MockWorkerToPrimary, WorkerToPrimary, WorkerToPrimaryServer},
     worker_to_worker_client::WorkerToWorkerClient,
-    worker_to_worker_server::{WorkerToWorker, WorkerToWorkerServer},
+    worker_to_worker_server::{MockWorkerToWorker, WorkerToWorker, WorkerToWorkerServer},
     CertificateDigest as CertificateDigestProto, Collection, CollectionError,
     CollectionRetrievalResult, Empty, GetCollectionsRequest, GetCollectionsResponse,
     GetPrimaryAddressResponse, MultiAddr as MultiAddrProto, NewEpochRequest, NewNetworkInfoRequest,