Introduce the BeaconStore

Alenar · Alenar · commit 3ecd83489683 · 2022-05-13T13:53:36.000+02:00
diff --git a/mithril-aggregator/src/beacon_store.rs b/mithril-aggregator/src/beacon_store.rs
@@ -0,0 +1,82 @@
+use async_trait::async_trait;
+use mithril_common::entities::Beacon;
+
+#[cfg(test)]
+use mockall::automock;
+
+/// BeaconStore represents a beacon store interactor
+#[cfg_attr(test, automock)]
+#[async_trait]
+pub trait BeaconStore: Sync + Send {
+    /// Get the current beacon
+    async fn get_current_beacon(&self) -> Result<Option<Beacon>, String>;
+
+    /// Set the current beacon
+    async fn set_current_beacon(&mut self, beacon: Beacon) -> Result<(), String>;
+
+    /// Reset the current beacon
+    async fn reset_current_beacon(&mut self) -> Result<(), String>;
+}
+
+/// MemoryBeaconStore is in memory [`BeaconStore`]
+pub struct MemoryBeaconStore {
+    current_beacon: Option<Beacon>,
+}
+
+impl Default for MemoryBeaconStore {
+    /// MemoryBeaconStore factory
+    fn default() -> Self {
+        Self {
+            current_beacon: None,
+        }
+    }
+}
+
+#[async_trait]
+impl BeaconStore for MemoryBeaconStore {
+    async fn get_current_beacon(&self) -> Result<Option<Beacon>, String> {
+        Ok(self.current_beacon.clone())
+    }
+
+    async fn set_current_beacon(&mut self, beacon: Beacon) -> Result<(), String> {
+        self.current_beacon = Some(beacon);
+        Ok(())
+    }
+
+    async fn reset_current_beacon(&mut self) -> Result<(), String> {
+        self.current_beacon = None;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{BeaconStore, MemoryBeaconStore};
+    use mithril_common::fake_data;
+
+    #[tokio::test]
+    async fn test_can_store_beacon() {
+        let mut sut = MemoryBeaconStore::default();
+        let beacon = fake_data::beacon();
+        sut.set_current_beacon(beacon.clone())
+            .await
+            .expect("unexpected error in set_current_beacon");
+        let stored_beacon = sut.get_current_beacon().await;
+
+        assert_eq!(Some(beacon), stored_beacon.unwrap());
+    }
+
+    #[tokio::test]
+    async fn test_reset_current_beacon_ok() {
+        let mut sut = MemoryBeaconStore::default();
+        sut.set_current_beacon(fake_data::beacon())
+            .await
+            .expect("unexpected error in set_current_beacon");
+        sut.reset_current_beacon()
+            .await
+            .expect("unexpected error in set_current_beacon");
+        let stored_beacon = sut.get_current_beacon().await;
+
+        assert_eq!(None, stored_beacon.unwrap());
+    }
+}
diff --git a/mithril-aggregator/src/dependency.rs b/mithril-aggregator/src/dependency.rs
@@ -1,10 +1,14 @@
+use crate::beacon_store::BeaconStore;
 use std::sync::Arc;
 use tokio::sync::RwLock;
 
 use super::entities::*;
 use super::multi_signer::MultiSigner;
 use super::snapshot_store::SnapshotStorer;
 
+/// BeaconStoreWrapper wraps a BeaconStore
+pub type BeaconStoreWrapper = Arc<RwLock<dyn BeaconStore>>;
+
 ///  SnapshotStorerWrapper wraps a SnapshotStorer
 pub type SnapshotStorerWrapper = Arc<RwLock<dyn SnapshotStorer>>;
 
@@ -16,6 +20,7 @@ pub struct DependencyManager {
     pub config: Config,
     pub snapshot_storer: Option<SnapshotStorerWrapper>,
     pub multi_signer: Option<MultiSignerWrapper>,
+    pub beacon_store: Option<BeaconStoreWrapper>,
 }
 
 impl DependencyManager {
@@ -25,6 +30,7 @@ impl DependencyManager {
             config,
             snapshot_storer: None,
             multi_signer: None,
+            beacon_store: None,
         }
     }
 
@@ -39,4 +45,10 @@ impl DependencyManager {
         self.multi_signer = Some(multi_signer);
         self
     }
+
+    /// With MultiSigner
+    pub fn with_beacon_store(&mut self, beacon_store: BeaconStoreWrapper) -> &mut Self {
+        self.beacon_store = Some(beacon_store);
+        self
+    }
 }
diff --git a/mithril-aggregator/src/http_server.rs b/mithril-aggregator/src/http_server.rs
@@ -10,7 +10,9 @@ use std::sync::Arc;
 use warp::Future;
 use warp::{http::Method, http::StatusCode, Filter};
 
-use super::dependency::{DependencyManager, MultiSignerWrapper, SnapshotStorerWrapper};
+use super::dependency::{
+    BeaconStoreWrapper, DependencyManager, MultiSignerWrapper, SnapshotStorerWrapper,
+};
 use super::multi_signer;
 
 const SERVER_BASE_PATH: &str = "aggregator";
@@ -67,10 +69,11 @@ mod router {
 
     /// GET /certificate-pending
     pub fn certificate_pending(
-        _dependency_manager: Arc<DependencyManager>,
+        dependency_manager: Arc<DependencyManager>,
     ) -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
         warp::path!("certificate-pending")
             .and(warp::get())
+            .and(with_beacon_store(dependency_manager))
             .and_then(handlers::certificate_pending)
     }
 
@@ -138,19 +141,44 @@ mod router {
     ) -> impl Filter<Extract = (MultiSignerWrapper,), Error = Infallible> + Clone {
         warp::any().map(move || dependency_manager.multi_signer.as_ref().unwrap().clone())
     }
+
+    /// With beacon store middleware
+    fn with_beacon_store(
+        dependency_manager: Arc<DependencyManager>,
+    ) -> impl Filter<Extract = (BeaconStoreWrapper,), Error = Infallible> + Clone {
+        warp::any().map(move || dependency_manager.beacon_store.as_ref().unwrap().clone())
+    }
 }
 
 mod handlers {
     use super::*;
 
     /// Certificate Pending
-    pub async fn certificate_pending() -> Result<impl warp::Reply, Infallible> {
+    pub async fn certificate_pending(
+        beacon_store: BeaconStoreWrapper,
+    ) -> Result<impl warp::Reply, Infallible> {
         debug!("certificate_pending");
 
-        // Certificate pending
-        let certificate_pending = fake_data::certificate_pending();
+        let beacon_store = beacon_store.read().await;
+        match beacon_store.get_current_beacon().await {
+            Ok(Some(beacon)) => {
+                let mut certificate_pending = fake_data::certificate_pending();
+                certificate_pending.beacon = beacon;
 
-        Ok(warp::reply::json(&certificate_pending))
+                Ok(warp::reply::with_status(
+                    warp::reply::json(&certificate_pending),
+                    StatusCode::OK,
+                ))
+            }
+            Ok(None) => Ok(warp::reply::with_status(
+                warp::reply::json(&Null),
+                StatusCode::NO_CONTENT,
+            )),
+            Err(err) => Ok(warp::reply::with_status(
+                warp::reply::json(&entities::Error::new("MITHRIL-E0006".to_string(), err)),
+                StatusCode::INTERNAL_SERVER_ERROR,
+            )),
+        }
     }
 
     /// Certificate by certificate hash
@@ -329,6 +357,7 @@ mod tests {
     use tokio::sync::RwLock;
     use warp::test::request;
 
+    use super::super::beacon_store::MockBeaconStore;
     use super::super::entities::*;
     use super::super::multi_signer::MockMultiSigner;
     use super::super::snapshot_store::MockSnapshotStorer;
@@ -345,9 +374,66 @@ mod tests {
 
     #[tokio::test]
     async fn test_certificate_pending_get_ok() {
-        let dependency_manager = setup_dependency_manager();
+        let mut beacon_store = MockBeaconStore::new();
+        beacon_store
+            .expect_get_current_beacon()
+            .return_once(|| Ok(Some(fake_data::beacon())));
+        let mut dependency_manager = setup_dependency_manager();
+        let method = Method::GET.as_str();
+        let path = "/certificate-pending";
+        dependency_manager.with_beacon_store(Arc::new(RwLock::new(beacon_store)));
+
+        let response = request()
+            .method(method)
+            .path(&format!("/{}{}", SERVER_BASE_PATH, path))
+            .reply(&router::routes(Arc::new(dependency_manager)))
+            .await;
+
+        APISpec::from_file(API_SPEC_FILE)
+            .method(method)
+            .path(path)
+            .validate_request(&Null)
+            .unwrap()
+            .validate_response(&response)
+            .expect("OpenAPI error");
+    }
+
+    #[tokio::test]
+    async fn test_certificate_pending_get_ok_204() {
+        let mut beacon_store = MockBeaconStore::new();
+        beacon_store
+            .expect_get_current_beacon()
+            .return_once(|| Ok(None));
+        let mut dependency_manager = setup_dependency_manager();
+        let method = Method::GET.as_str();
+        let path = "/certificate-pending";
+        dependency_manager.with_beacon_store(Arc::new(RwLock::new(beacon_store)));
+
+        let response = request()
+            .method(method)
+            .path(&format!("/{}{}", SERVER_BASE_PATH, path))
+            .reply(&router::routes(Arc::new(dependency_manager)))
+            .await;
+
+        APISpec::from_file(API_SPEC_FILE)
+            .method(method)
+            .path(path)
+            .validate_request(&Null)
+            .unwrap()
+            .validate_response(&response)
+            .expect("OpenAPI error");
+    }
+
+    #[tokio::test]
+    async fn test_certificate_pending_get_ko() {
+        let mut beacon_store = MockBeaconStore::new();
+        beacon_store
+            .expect_get_current_beacon()
+            .return_once(|| Err("an error occurred".to_string()));
+        let mut dependency_manager = setup_dependency_manager();
         let method = Method::GET.as_str();
         let path = "/certificate-pending";
+        dependency_manager.with_beacon_store(Arc::new(RwLock::new(beacon_store)));
 
         let response = request()
             .method(method)
diff --git a/mithril-aggregator/src/lib.rs b/mithril-aggregator/src/lib.rs
@@ -1,3 +1,4 @@
+mod beacon_store;
 mod dependency;
 mod entities;
 mod http_server;
@@ -13,6 +14,7 @@ pub use crate::multi_signer::{
     ProtocolSignerVerificationKey, ProtocolStake,
 };
 pub use crate::snapshot_store::SnapshotStoreHTTPClient;
+pub use beacon_store::MemoryBeaconStore;
 pub use dependency::DependencyManager;
 pub use runtime::AggregatorRuntime;
 pub use snapshotter::Snapshotter;
diff --git a/mithril-aggregator/src/main.rs b/mithril-aggregator/src/main.rs
@@ -3,9 +3,9 @@
 use clap::Parser;
 
 use mithril_aggregator::{
-    key_decode_hex, AggregatorRuntime, Config, DependencyManager, MultiSigner, MultiSignerImpl,
-    ProtocolParameters, ProtocolPartyId, ProtocolSignerVerificationKey, ProtocolStake, Server,
-    SnapshotStoreHTTPClient,
+    key_decode_hex, AggregatorRuntime, Config, DependencyManager, MemoryBeaconStore, MultiSigner,
+    MultiSignerImpl, ProtocolParameters, ProtocolPartyId, ProtocolSignerVerificationKey,
+    ProtocolStake, Server, SnapshotStoreHTTPClient,
 };
 use mithril_common::fake_data;
 use slog::{Drain, Level, Logger};
@@ -88,18 +88,23 @@ async fn main() {
     )));
 
     let multi_signer = Arc::new(RwLock::new(init_multi_signer()));
+    let beacon_store = Arc::new(RwLock::new(MemoryBeaconStore::default()));
 
-    // Init dependecy manager
+    // Init dependency manager
     let mut dependency_manager = DependencyManager::new(config);
     dependency_manager
         .with_snapshot_storer(snapshot_storer.clone())
-        .with_multi_signer(multi_signer.clone());
+        .with_multi_signer(multi_signer.clone())
+        .with_beacon_store(beacon_store.clone());
     let dependency_manager = Arc::new(dependency_manager);
 
     // Start snapshot uploader
     let handle = tokio::spawn(async move {
-        let runtime =
-            AggregatorRuntime::new(args.snapshot_interval * 1000, args.db_directory.clone());
+        let runtime = AggregatorRuntime::new(
+            args.snapshot_interval * 1000,
+            args.db_directory.clone(),
+            beacon_store.clone(),
+        );
         runtime.run().await
     });
 
diff --git a/mithril-aggregator/src/runtime.rs b/mithril-aggregator/src/runtime.rs
