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 std::sync::Arc; 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; use surrealdb::engine::remote::ws::Client; use surrealdb::Surreal; 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() -> Surreal { let db = surreal_brain::db::db_connection(DB_URL, DB_NS, DB_NAME).await.unwrap(); DB_STATE .get_or_init(|| async { let old_brain_data = surreal_brain::old_brain::BrainData::load_from_disk().unwrap(); db.query(format!("REMOVE DATABASE {DB_NAME}")).await.unwrap(); db.query(std::fs::read_to_string("interim_tables.surql").unwrap()).await.unwrap(); surreal_brain::db::migration0(&db, &old_brain_data).await.unwrap(); }) .await; db } 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 { let db = surreal_brain::db::db_connection(DB_URL, DB_NS, DB_NAME).await.unwrap(); let db_arc = Arc::new(db); Server::builder() .add_service(BrainGeneralCliServer::new(BrainGeneralCliForReal::new(db_arc.clone()))) .add_service(BrainVmCliServer::new(BrainVmCliForReal::new(db_arc.clone()))) .add_service(BrainVmDaemonServer::new(BrainVmDaemonForReal::new(db_arc.clone()))) .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 { let db = surreal_brain::db::db_connection(DB_URL, DB_NS, DB_NAME).await.unwrap(); let db_arc = Arc::new(db); tonic::transport::Server::builder() .add_service(BrainGeneralCliServer::new(BrainGeneralCliForReal::new(db_arc.clone()))) .add_service(BrainVmCliServer::new(BrainVmCliForReal::new(db_arc.clone()))) .add_service(BrainVmDaemonServer::new(BrainVmDaemonForReal::new(db_arc.clone()))) .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 }