use detee_shared::general_proto::brain_general_cli_server::BrainGeneralCliServer; use detee_shared::vm_proto::brain_vm_cli_server::BrainVmCliServer; use detee_shared::vm_proto::brain_vm_daemon_server::BrainVmDaemonServer; use dotenv::dotenv; use hyper_util::rt::TokioIo; use std::net::SocketAddr; use std::sync::Arc; use surreal_brain::grpc; use surrealdb::engine::remote::ws::Client; use surrealdb::Surreal; use tokio::io::DuplexStream; use tokio::net::TcpListener; use tokio::sync::OnceCell; use tonic::transport::{Channel, Endpoint, Server, Uri}; use tower::service_fn; pub static DB_STATE: OnceCell<()> = OnceCell::const_new(); pub async fn prepare_test_db() -> Surreal { dotenv().ok(); let db_url = std::env::var("DB_URL").expect("DB_URL not set in .env"); let db_user = std::env::var("DB_USER").expect("DB_USER not set in .env"); let db_pass = std::env::var("DB_PASS").expect("DB_PASS not set in .env"); let db_ns = "test_brain"; let db_name = "test_migration_db"; let db = surreal_brain::db::db_connection(&db_url, &db_user, &db_pass, &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 { dotenv().ok(); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let db_url = std::env::var("DB_URL").expect("DB_URL not set in .env"); let db_user = std::env::var("DB_USER").expect("DB_USER not set in .env"); let db_pass = std::env::var("DB_PASS").expect("DB_PASS not set in .env"); let db_ns = "test_brain"; let db_name = "test_migration_db"; tokio::spawn(async move { let db = surreal_brain::db::db_connection(&db_url, &db_user, &db_pass, &db_ns, &db_name) .await .unwrap(); let db_arc = Arc::new(db); Server::builder() .add_service(BrainGeneralCliServer::new(grpc::general::BrainGeneralCliServer::new( db_arc.clone(), ))) .add_service(BrainVmCliServer::new(grpc::vm::BrainVmCliServer::new(db_arc.clone()))) .add_service(BrainVmDaemonServer::new(grpc::vm::BrainVmDaemonServer::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 { dotenv().ok(); let (client, server) = tokio::io::duplex(1024); let db_url = std::env::var("DB_URL").expect("DB_URL not set in .env"); let db_user = std::env::var("DB_USER").expect("DB_USER not set in .env"); let db_pass = std::env::var("DB_PASS").expect("DB_PASS not set in .env"); let db_ns = "test_brain"; let db_name = "test_migration_db"; tokio::spawn(async move { let db = surreal_brain::db::db_connection(&db_url, &db_user, &db_pass, &db_ns, &db_name) .await .unwrap(); let db_arc = Arc::new(db); tonic::transport::Server::builder() .add_service(BrainGeneralCliServer::new(grpc::general::BrainGeneralCliServer::new(db_arc.clone()))) .add_service(BrainVmCliServer::new(grpc::vm::BrainVmCliServer::new(db_arc.clone()))) .add_service(BrainVmDaemonServer::new(grpc::vm::BrainVmDaemonServer::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 }