From 93f063bb01c17e12f214c7cde069fa4e78ba6a78 Mon Sep 17 00:00:00 2001
From: mfahampshire <maxhampshire@pm.me>
Date: Fri, 3 Jan 2025 13:28:02 +0100
Subject: [PATCH] added test start fn to client_pool: uses specified gw for all
 clients

---
 .../src/client_pool/mixnet_client_pool.rs     | 47 +++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/sdk/rust/nym-sdk/src/client_pool/mixnet_client_pool.rs b/sdk/rust/nym-sdk/src/client_pool/mixnet_client_pool.rs
index 41656fe272..ff7b8b8198 100644
--- a/sdk/rust/nym-sdk/src/client_pool/mixnet_client_pool.rs
+++ b/sdk/rust/nym-sdk/src/client_pool/mixnet_client_pool.rs
@@ -1,5 +1,9 @@
+// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: GPL-3.0-only
+
 use crate::mixnet::{MixnetClient, MixnetClientBuilder, NymNetworkDetails};
 use anyhow::Result;
+use nym_crypto::asymmetric::ed25519;
 use std::fmt;
 use std::sync::Arc;
 use tokio::sync::RwLock;
@@ -108,6 +112,49 @@ impl ClientPool {
         }
     }
 
+    // Even though this is basically start() with an extra param since I think this
+    // will only be used for testing scenarios, and I didn't want to unnecessarily add
+    // another param to the function that will be used elsewhere, hence this is its own fn
+    pub async fn start_with_specified_gateway(&self, gateway: ed25519::PublicKey) -> Result<()> {
+        loop {
+            let spawned_clients = self.clients.read().await.len();
+            let addresses = self;
+            debug!(
+                "Currently spawned clients: {}: {:?}",
+                spawned_clients, addresses
+            );
+            if self.cancel_token.is_cancelled() {
+                break Ok(());
+            }
+            if spawned_clients >= self.client_pool_reserve_number {
+                debug!("Got enough clients already: sleeping");
+            } else {
+                info!(
+                    "Clients in reserve = {}, reserve amount = {}, spawning new client",
+                    spawned_clients, self.client_pool_reserve_number
+                );
+                let client = loop {
+                    let net = NymNetworkDetails::new_from_env();
+                    match MixnetClientBuilder::new_ephemeral()
+                        .network_details(net)
+                        .request_gateway(gateway.to_string())
+                        .build()?
+                        .connect_to_mixnet()
+                        .await
+                    {
+                        Ok(client) => break client,
+                        Err(err) => {
+                            warn!("Error creating client: {:?}, will retry in 100ms", err);
+                            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+                        }
+                    }
+                };
+                self.clients.write().await.push(Arc::new(client));
+            }
+            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+        }
+    }
+
     pub async fn disconnect_pool(&self) {
         info!("Triggering Client Pool disconnect");
         self.cancel_token.cancel();