From 5925b0ce1c2202cd05fe966894bc0c586371e820 Mon Sep 17 00:00:00 2001
From: Cyle Witruk <cyle.witruk@outlook.com>
Date: Thu, 3 Oct 2024 08:29:00 +0200
Subject: [PATCH] add .with_mocked_clients() convenience method to reduce
 boilerplate

---
 signer/src/api/new_block.rs           |  6 ++---
 signer/src/context/mod.rs             |  3 +--
 signer/src/network/libp2p/swarm.rs    |  3 +--
 signer/src/network/mod.rs             |  6 ++---
 signer/src/testing/context.rs         | 39 +++++++++++++++++++++------
 signer/src/transaction_coordinator.rs |  3 +--
 signer/tests/integration/postgres.rs  |  3 +--
 7 files changed, 39 insertions(+), 24 deletions(-)

diff --git a/signer/src/api/new_block.rs b/signer/src/api/new_block.rs
index b495ebce..99c0e755 100644
--- 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
index 8338d0ef..7bb5187f 100644
--- 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
index 1614228f..e23be9f2 100644
--- 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
index 9f8ee648..88fc9855 100644
--- 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
index d11ed815..32287f08 100644
--- 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
index 17007fca..2892dbdf 100644
--- 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
index 2f664bc4..448e4fc1 100644
--- 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 {