add .with_mocked_clients() convenience method to reduce boilerplate

stacks-network · Oct 3, 2024 · 5925b0c · 5925b0c
1 parent a5845af
commit 5925b0c
Show file tree

Hide file tree

Showing 7 changed files with 39 additions and 24 deletions.
diff --git a/signer/src/api/new_block.rs b/signer/src/api/new_block.rs
@@ -160,8 +160,7 @@ mod tests {
     {
         let ctx = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         let api = ApiState { ctx: ctx.clone() };
@@ -192,8 +191,7 @@ mod tests {
     {
         let ctx = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         let api = ApiState { ctx: ctx.clone() };

diff --git a/signer/src/context/mod.rs b/signer/src/context/mod.rs
@@ -185,8 +185,7 @@ mod tests {
         // Create a context.
         let context = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         // Clone the context.

diff --git a/signer/src/network/libp2p/swarm.rs b/signer/src/network/libp2p/swarm.rs
@@ -248,8 +248,7 @@ mod tests {
 
         let ctx = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         let term = ctx.get_termination_handle();

diff --git a/signer/src/network/mod.rs b/signer/src/network/mod.rs
@@ -143,14 +143,12 @@ mod tests {
 
         let context1 = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         let context2 = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         let term1 = context1.get_termination_handle();

diff --git a/signer/src/testing/context.rs b/signer/src/testing/context.rs
@@ -509,6 +509,34 @@ where
 {
 }
 
+/// Trait for configuring the context with mocked clients. These methods are
+/// available when no clients have been configured yet.
+pub trait ConfigureMockedClients<Storage>
+where
+    Self: Sized + BuilderState<Storage, (), ()>,
+{
+    /// Configure the context to use mocks for all client implementations.
+    fn with_mocked_clients(
+        self,
+    ) -> ContextBuilder<Storage, WrappedMock<MockBitcoinInteract>, WrappedMock<MockStacksInteract>>
+    {
+        let config = self.get_config();
+        ContextBuilder {
+            config: ContextConfig {
+                settings: config.settings,
+                storage: config.storage,
+                bitcoin: WrappedMock::default(),
+                stacks: WrappedMock::default(),
+            },
+        }
+    }
+}
+
+impl<Storage> ConfigureMockedClients<Storage> for ContextBuilder<Storage, (), ()> where
+    Self: Sized + BuilderState<Storage, (), ()>
+{
+}
+
 /// Trait for building a [`TestContext`]. The [`BuildContext::build`] method
 /// consumes the builder and returns a new [`TestContext`]. The method is only
 /// available when all required components have been configured.
@@ -555,18 +583,14 @@ mod tests {
 
     use crate::{
         context::{Context as _, SignerEvent, SignerSignal},
-        testing::context::{
-            BuildContext, ConfigureBitcoinClient, ConfigureStacksClient, ConfigureStorage,
-            ContextBuilder, TestContext,
-        },
+        testing::context::*,
     };
 
     #[test]
     fn can_build() {
         let _builder = ContextBuilder::new()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
     }
 
@@ -576,8 +600,7 @@ mod tests {
     async fn context_clone_signalling_works() {
         let context = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         let context = Arc::new(context);

diff --git a/signer/src/transaction_coordinator.rs b/signer/src/transaction_coordinator.rs
@@ -617,8 +617,7 @@ mod tests {
 
         let context = TestContext::builder()
             .with_in_memory_storage()
-            .with_mocked_bitcoin_client()
-            .with_mocked_stacks_client()
+            .with_mocked_clients()
             .build();
 
         testing::transaction_coordinator::TestEnvironment {

diff --git a/signer/tests/integration/postgres.rs b/signer/tests/integration/postgres.rs
@@ -1384,8 +1384,7 @@ async fn transaction_coordinator_test_environment(
 
     let context = TestContext::builder()
         .with_storage(store)
-        .with_mocked_bitcoin_client()
-        .with_mocked_stacks_client()
+        .with_mocked_clients()
         .build();
 
     testing::transaction_coordinator::TestEnvironment {