From b241765b3aea0abbe9f4f60b58696229fb243629 Mon Sep 17 00:00:00 2001
From: Dafna Matsry <dafna.matsry@starkware.co>
Date: Thu, 5 Dec 2024 13:57:40 +0200
Subject: [PATCH] refactor(starknet_batcher): add getters to the proposal
 manger

---
 crates/starknet_batcher/src/batcher_test.rs   | 22 ++++++++++
 .../starknet_batcher/src/proposal_manager.rs  | 40 ++++++++++++++-----
 2 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/crates/starknet_batcher/src/batcher_test.rs b/crates/starknet_batcher/src/batcher_test.rs
index e6ca53d1e1..05ab9c25ae 100644
--- a/crates/starknet_batcher/src/batcher_test.rs
+++ b/crates/starknet_batcher/src/batcher_test.rs
@@ -580,6 +580,14 @@ trait ProposalManagerTraitWrapper: Send + Sync {
         proposal_id: ProposalId,
     ) -> BoxFuture<'_, ProposalResult<ProposalOutput>>;
 
+    fn wrap_get_active_proposal(&self) -> BoxFuture<'_, Option<ProposalId>>;
+
+    fn wrap_get_completed_proposals(
+        &self,
+    ) -> BoxFuture<'_, Arc<tokio::sync::Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>>>;
+
+    fn wrap_await_active_proposal(&mut self) -> BoxFuture<'_, bool>;
+
     fn wrap_get_proposal_status(
         &self,
         proposal_id: ProposalId,
@@ -613,6 +621,20 @@ impl<T: ProposalManagerTraitWrapper> ProposalManagerTrait for T {
         self.wrap_take_proposal_result(proposal_id).await
     }
 
+    async fn get_active_proposal(&self) -> Option<ProposalId> {
+        self.wrap_get_active_proposal().await
+    }
+
+    async fn get_completed_proposals(
+        &self,
+    ) -> Arc<tokio::sync::Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>> {
+        self.wrap_get_completed_proposals().await
+    }
+
+    async fn await_active_proposal(&mut self) -> bool {
+        self.wrap_await_active_proposal().await
+    }
+
     async fn get_proposal_status(&self, proposal_id: ProposalId) -> InternalProposalStatus {
         self.wrap_get_proposal_status(proposal_id).await
     }
diff --git a/crates/starknet_batcher/src/proposal_manager.rs b/crates/starknet_batcher/src/proposal_manager.rs
index d588861c10..b910e46280 100644
--- a/crates/starknet_batcher/src/proposal_manager.rs
+++ b/crates/starknet_batcher/src/proposal_manager.rs
@@ -63,6 +63,16 @@ pub trait ProposalManagerTrait: Send + Sync {
         proposal_id: ProposalId,
     ) -> ProposalResult<ProposalOutput>;
 
+    #[allow(dead_code)]
+    async fn get_active_proposal(&self) -> Option<ProposalId>;
+
+    #[allow(dead_code)]
+    async fn get_completed_proposals(
+        &self,
+    ) -> Arc<Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>>;
+
+    async fn await_active_proposal(&mut self) -> bool;
+
     async fn get_proposal_status(&self, proposal_id: ProposalId) -> InternalProposalStatus;
 
     async fn await_proposal_commitment(
@@ -162,6 +172,26 @@ impl ProposalManagerTrait for ProposalManager {
             .ok_or(GetProposalResultError::ProposalDoesNotExist { proposal_id })?
     }
 
+    async fn get_active_proposal(&self) -> Option<ProposalId> {
+        *self.active_proposal.lock().await
+    }
+
+    async fn get_completed_proposals(
+        &self,
+    ) -> Arc<Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>> {
+        self.executed_proposals.clone()
+    }
+
+    // Awaits the active proposal.
+    // Returns true if there was an active proposal, and false otherwise.
+    async fn await_active_proposal(&mut self) -> bool {
+        if let Some(proposal_task) = self.active_proposal_task.take() {
+            proposal_task.join_handle.await.ok();
+            return true;
+        }
+        false
+    }
+
     // Returns None if the proposal does not exist, otherwise, returns the status of the proposal.
     async fn get_proposal_status(&self, proposal_id: ProposalId) -> InternalProposalStatus {
         match self.executed_proposals.lock().await.get(&proposal_id) {
@@ -245,16 +275,6 @@ impl ProposalManager {
         Ok(())
     }
 
-    // Awaits the active proposal.
-    // Returns true if there was an active proposal, and false otherwise.
-    pub async fn await_active_proposal(&mut self) -> bool {
-        if let Some(proposal_task) = self.active_proposal_task.take() {
-            proposal_task.join_handle.await.ok();
-            return true;
-        }
-        false
-    }
-
     // Ends the current active proposal.
     // This call is non-blocking.
     async fn abort_active_proposal(&mut self) {