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 { let (client, server) = tokio::io::duplex(1024); 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 } pub async fn connect_stream_client_channel(c_stream: DuplexStream) -> Channel { let mut client = Some(c_stream); Endpoint::from_static("http://127.0.0.1:0") .connect_with_connector(service_fn(move |_: Uri| { let client = client.take().unwrap(); async move { Ok::, std::io::Error>(TokioIo::new(client)) } })) .await .unwrap() } pub async fn run_service_for_stream() -> Channel { let client = run_service_for_stream_server().await; connect_stream_client_channel(client).await }