testnet/brain
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor mock vm daemon and simplified client service fn

Browse Source
This commit is contained in:
Noor 2025-04-30 16:36:25 +05:30
parent ab39b12da9
commit de5fba37c9
Signed by: noormohammedb
GPG Key ID: D83EFB8B3B967146
3 changed files with 15 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
tests/common/prepare_test_env.rs
View File

@ -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
}

11
tests/common/vm_daemon_utils.rs
View File

@ -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(

10
tests/grpc_vm_daemon_test.rs
View File

@ -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();