From de5fba37c9bc47a4ddaaa74aebdf82b542d8414d Mon Sep 17 00:00:00 2001
From: Noor <noormohammedb@protonmail.com>
Date: Wed, 30 Apr 2025 16:36:25 +0530
Subject: [PATCH] refactor mock vm daemon and simplified client service fn

---
 tests/common/prepare_test_env.rs | 28 ++++++++--------------------
 tests/common/vm_daemon_utils.rs  | 11 +++--------
 tests/grpc_vm_daemon_test.rs     | 10 ++++------
 3 files changed, 15 insertions(+), 34 deletions(-)

diff --git a/tests/common/prepare_test_env.rs b/tests/common/prepare_test_env.rs
index 0fd6699..cdeaed8 100644
--- a/tests/common/prepare_test_env.rs
+++ b/tests/common/prepare_test_env.rs
@@ -54,12 +54,9 @@ pub async fn run_service_in_background() -> SocketAddr {
     addr
 }
 
-pub async fn run_service_for_stream_server() -> (DuplexStream, SocketAddr) {
+pub async fn run_service_for_stream_server() -> DuplexStream {
     let (client, server) = tokio::io::duplex(1024);
 
-    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
-    let addr = listener.local_addr().unwrap();
-
     tokio::spawn(async move {
         tonic::transport::Server::builder()
             .add_service(BrainGeneralCliServer::new(BrainGeneralCliForReal {}))
@@ -68,31 +65,22 @@ pub async fn run_service_for_stream_server() -> (DuplexStream, SocketAddr) {
             .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
             .await
     });
-    (client, addr)
+    client
 }
 
-pub async fn connect_stream_client_channel(c_stream: DuplexStream, addr: SocketAddr) -> Channel {
-    let url = format!("http://{}", addr);
+pub async fn connect_stream_client_channel(c_stream: DuplexStream) -> Channel {
     let mut client = Some(c_stream);
 
-    Endpoint::try_from(url)
-        .unwrap()
+    Endpoint::from_static("http://127.0.0.1:0")
         .connect_with_connector(service_fn(move |_: Uri| {
-            let client = client.take();
-
-            async move {
-                if let Some(client) = client {
-                    Ok(TokioIo::new(client))
-                } else {
-                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Client already taken"))
-                }
-            }
+            let client = client.take().unwrap();
+            async move { Ok::<TokioIo<DuplexStream>, std::io::Error>(TokioIo::new(client)) }
         }))
         .await
         .unwrap()
 }
 
 pub async fn run_service_for_stream() -> Channel {
-    let (client, addr) = run_service_for_stream_server().await;
-    connect_stream_client_channel(client, addr).await
+    let client = run_service_for_stream_server().await;
+    connect_stream_client_channel(client).await
 }
diff --git a/tests/common/vm_daemon_utils.rs b/tests/common/vm_daemon_utils.rs
index 313dd44..ac41f59 100644
--- a/tests/common/vm_daemon_utils.rs
+++ b/tests/common/vm_daemon_utils.rs
@@ -1,4 +1,3 @@
-use super::prepare_test_env::{connect_stream_client_channel, run_service_for_stream_server};
 use super::test_utils::Key;
 use detee_shared::vm_proto;
 use detee_shared::vm_proto::brain_vm_daemon_client::BrainVmDaemonClient;
@@ -8,12 +7,8 @@ use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::transport::Channel;
 
-pub async fn mock_vm_daemon() -> (Channel, String) {
-    let (tokio_duplex, addr) = run_service_for_stream_server().await;
-
-    let channel = connect_stream_client_channel(tokio_duplex, addr).await;
-
-    let daemon_client = BrainVmDaemonClient::new(channel.clone());
+pub async fn mock_vm_daemon(brain_channel: Channel) -> String {
+    let daemon_client = BrainVmDaemonClient::new(brain_channel);
     let daemon_key = Key::new();
 
     register_vm_node(daemon_client.clone(), daemon_key.clone(), Key::new().pubkey).await;
@@ -34,7 +29,7 @@ pub async fn mock_vm_daemon() -> (Channel, String) {
 
     tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
 
-    (channel, daemon_key.pubkey)
+    daemon_key.pubkey
 }
 
 pub async fn register_vm_node(
diff --git a/tests/grpc_vm_daemon_test.rs b/tests/grpc_vm_daemon_test.rs
index b10ca18..681ae13 100644
--- a/tests/grpc_vm_daemon_test.rs
+++ b/tests/grpc_vm_daemon_test.rs
@@ -1,5 +1,5 @@
 use common::{
-    prepare_test_env::{prepare_test_db, run_service_in_background},
+    prepare_test_env::{prepare_test_db, run_service_for_stream, run_service_in_background},
     test_utils::Key,
     vm_daemon_utils::{mock_vm_daemon, register_vm_node},
 };
@@ -30,11 +30,9 @@ async fn test_brain_message() {
     env_logger::builder().filter_level(log::LevelFilter::Info).init();
     let _ = prepare_test_db().await;
 
-    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
-
-    let (channel, daemon_key) = mock_vm_daemon().await;
-
-    let mut cli_client = BrainVmCliClient::new(channel);
+    let brain_channel = run_service_for_stream().await;
+    let daemon_key = mock_vm_daemon(brain_channel.clone()).await;
+    let mut cli_client = BrainVmCliClient::new(brain_channel);
 
     let cli_key = Key::new();