WIP: test coordinator::core::Service

rust/Cargo.toml

@@ -37,6 +37,9 @@ opentelemetry = { version = "0.2.0", optional = true }
 tracing-opentelemetry = { version = "0.2.0", optional = true }
 opentelemetry-jaeger = { version = "0.1.0", optional = true }
 
+[dev-dependencies]
+mockall = "0.6.0"
+
 [[bin]]
 name = "coordinator"
 path = "src/bin/coordinator.rs"

rust/src/aggregator/mod.rs

@@ -1,5 +1,3 @@
-#![allow(dead_code)]
-
 pub mod api;
 pub mod py_aggregator;
 pub mod rpc;

rust/src/aggregator/py_aggregator.rs

@@ -190,7 +190,7 @@ pub struct PyAggregatorHandle {
 
 impl Aggregator for PyAggregatorHandle {
     type Error = ();
-    type AggregateFut = Pin<Box<dyn Future<Output = Result<Bytes, ()>>>>;
+    type AggregateFut = Pin<Box<dyn Future<Output = Result<Bytes, ()>> + Send>>;
     type AddWeightsFut = Pin<Box<dyn Future<Output = Result<(), ()>> + Send>>;
 
     fn add_weights(&mut self, weights: Bytes) -> Self::AddWeightsFut {

rust/src/aggregator/rpc.rs

@@ -28,7 +28,12 @@ mod inner {
     }
 }
 
-pub use inner::{Rpc, RpcClient as Client};
+pub use inner::Rpc;
+
+#[cfg(test)]
+pub use crate::tests::mocks::rpc::aggregator::Client;
+#[cfg(not(test))]
+pub use inner::RpcClient as Client;
 
 /// A server that serves a single client. A new `Server` is created
 /// for each new client.

rust/src/coordinator/core/mod.rs

@@ -3,4 +3,6 @@ mod heartbeat;
 mod protocol;
 mod service;
 
+#[cfg(test)]
+pub(crate) use self::service::ServiceRequests;
 pub use self::service::{RequestError, Selector, Service, ServiceHandle};

rust/src/coordinator/core/service.rs

@@ -31,7 +31,7 @@ use tokio::{
     },
 };
 
-struct AggregationFuture(Pin<Box<dyn Future<Output = Result<(), ()>>>>);
+struct AggregationFuture(Pin<Box<dyn Future<Output = Result<(), ()>> + Send>>);
 
 impl Future for AggregationFuture {
     type Output = Result<(), ()>;
@@ -478,6 +478,7 @@ where
     }
 }
 
+#[derive(Debug)]
 pub struct RequestError;
 
 pub struct ServiceRequests(Pin<Box<dyn Stream<Item = Request> + Send>>);

rust/src/coordinator/rpc.rs

@@ -17,7 +17,41 @@ mod inner {
         async fn end_training(id: ClientId, success: bool);
     }
 }
-pub use inner::{Rpc, RpcClient as Client};
+pub use inner::Rpc;
+
+#[cfg(test)]
+mod mocks {
+    pub use super::*;
+
+    use futures::future;
+    use mockall::mock;
+    use std::io;
+    use tarpc::{client::Config, context::Context, rpc::Transport};
+
+    mock! {
+        pub NewClient {
+            fn spawn(self) -> io::Result<MockRpcClient>;
+        }
+    }
+
+    mock! {
+        pub RpcClient {
+            fn new<T: Transport<(), ()> + 'static>(config: Config, transport: T) -> MockNewClient;
+            fn end_training(&mut self, ctx: Context, id: ClientId, success: bool) -> future::Ready<io::Result<()>>;
+        }
+    }
+
+    impl Clone for MockRpcClient {
+        fn clone(&self) -> Self {
+            Self::default()
+        }
+    }
+}
+
+#[cfg(not(test))]
+pub use inner::RpcClient as Client;
+#[cfg(test)]
+pub use mocks::MockRpcClient as Client;
 
 impl Rpc for Server {
     type EndTrainingFut = Pin<Box<dyn Future<Output = ()> + Send>>;

rust/src/lib.rs

@@ -10,3 +10,6 @@ extern crate serde;
 pub mod aggregator;
 pub mod common;
 pub mod coordinator;
+
+#[cfg(test)]
+mod tests;

rust/src/tests/coordinator.rs

@@ -0,0 +1,77 @@
+use crate::{
+    coordinator::{core::Service, models::HeartBeatResponse, settings::FederatedLearningSettings},
+    tests::mocks::{
+        coordinator::{MaxSelector, ServiceHandle},
+        rpc::aggregator::{Client, MockClient},
+    },
+};
+use futures::future;
+use tokio::time::{delay_for, Duration};
+
+#[cfg(logging)]
+use crate::common::{logging, settings::LoggingSettings};
+#[cfg(logging)]
+use tracing_subscriber::filter::EnvFilter;
+
+#[tokio::test]
+async fn test_rendez_vous_accept() {
+    #[cfg(logging)]
+    logging::configure(LoggingSettings {
+        telemetry: None,
+        filter: EnvFilter::try_new("trace").unwrap(),
+    });
+
+    let rpc_client: Client = MockClient::default().into();
+    let aggregator_url = "http://localhost:8082".to_string();
+
+    let (service_handle, service_requests) = ServiceHandle::new();
+
+    let service = Service::new(
+        MaxSelector,
+        FederatedLearningSettings {
+            rounds: 1,
+            participants_ratio: 1.0,
+            min_clients: 1,
+            heartbeat_timeout: 10,
+        },
+        aggregator_url.clone(),
+        rpc_client.clone(),
+        service_requests,
+    );
+    let _join_handle = tokio::spawn(service);
+
+    let id = service_handle.rendez_vous_accepted().await;
+    let round = service_handle.heartbeat_selected(id).await;
+    assert_eq!(round, 0);
+
+    rpc_client
+        .mock()
+        .expect_select()
+        .returning(|_, _| future::ready(Ok(Ok(()))));
+
+    let (url, _token) = service_handle.start_training_accepted(id).await;
+    assert_eq!(url, aggregator_url);
+
+    // pretend the client trained and sent its weights to the
+    // aggregator. The aggregator now sends an end training requests
+    // to the coordinator RPC server that we fake with the
+    // service_handle. The service should then trigger the aggregation
+    // and reject subsequent heartbeats and rendez-vous
+    rpc_client
+        .mock()
+        .expect_aggregate()
+        .returning(|_| future::ready(Ok(Ok(()))));
+
+    service_handle.end_training(id, true).await;
+    loop {
+        match service_handle.heartbeat(id).await {
+            HeartBeatResponse::StandBy => sleep_ms(10).await,
+            HeartBeatResponse::Finish => break,
+            _ => panic!("expected StandBy or Finish"),
+        }
+    }
+}
+
+async fn sleep_ms(ms: u64) {
+    delay_for(Duration::from_millis(ms)).await
+}