brain/tests/common/prepare_test_env.rs
Noor 32f6548eff
seperating tests
renamed a test file
seperated daemon service tests
modularizing test into reusable methods
2025-04-30 02:09:14 +05:30

100 lines
3.5 KiB
Rust

use detee_shared::vm_proto::brain_vm_cli_server::BrainVmCliServer;
use detee_shared::{
general_proto::brain_general_cli_server::BrainGeneralCliServer,
vm_proto::brain_vm_daemon_server::BrainVmDaemonServer,
};
use hyper_util::rt::TokioIo;
use std::net::SocketAddr;
use surreal_brain::grpc::{BrainGeneralCliForReal, BrainVmCliForReal, BrainVmDaemonForReal};
use tokio::io::DuplexStream;
use tokio::{net::TcpListener, sync::OnceCell};
use tonic::transport::{Channel, Endpoint, Server, Uri};
use tower::service_fn;
pub const DB_URL: &str = "localhost:8000";
pub const DB_NS: &str = "test_brain";
pub const DB_NAME: &str = "test_migration_db";
pub static DB_STATE: OnceCell<()> = OnceCell::const_new();
pub async fn prepare_test_db() {
DB_STATE
.get_or_init(|| async {
surreal_brain::db::init(DB_URL, DB_NS, DB_NAME)
.await
.expect("Failed to initialize the database");
let old_brain_data = surreal_brain::old_brain::BrainData::load_from_disk().unwrap();
surreal_brain::db::DB.query(format!("REMOVE DATABASE {DB_NAME}")).await.unwrap();
surreal_brain::db::DB
.query(std::fs::read_to_string("interim_tables.surql").unwrap())
.await
.unwrap();
surreal_brain::db::migration0(&old_brain_data).await.unwrap();
})
.await;
}
pub async fn run_service_in_background() -> SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
Server::builder()
.add_service(BrainGeneralCliServer::new(BrainGeneralCliForReal {}))
.add_service(BrainVmCliServer::new(BrainVmCliForReal {}))
.add_service(BrainVmDaemonServer::new(BrainVmDaemonForReal {}))
.serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
.await
.unwrap();
});
tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;
addr
}
pub async fn run_service_for_stream_server() -> (DuplexStream, SocketAddr) {
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 {}))
.add_service(BrainVmCliServer::new(BrainVmCliForReal {}))
.add_service(BrainVmDaemonServer::new(BrainVmDaemonForReal {}))
.serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
.await
});
(client, addr)
}
pub async fn connect_stream_client_channel(c_stream: DuplexStream, addr: SocketAddr) -> Channel {
let url = format!("http://{}", addr);
let mut client = Some(c_stream);
Endpoint::try_from(url)
.unwrap()
.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"))
}
}
}))
.await
.unwrap()
}
#[allow(dead_code)]
pub async fn run_service_for_stream() -> Channel {
let (client, addr) = run_service_for_stream_server().await;
connect_stream_client_channel(client, addr).await
}