brain/tests/common/prepare_test_env.rs

use anyhow::Result;
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::constants::DB_SCHEMA_FILES;
use surreal_brain::grpc::general::GeneralCliServer;
use surreal_brain::grpc::vm::{VmCliServer, VmDaemonServer};
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<Result<()>> = OnceCell::const_new();

pub async fn prepare_test_db() -> Result<Surreal<Client>> {
    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?;
    DB_STATE
        .get_or_init(|| async {
            let raw_mock_data = std::fs::read_to_string("tests/mock_data.yaml")?;
            let mock_data: surreal_brain::old_brain::BrainData =
                serde_yaml::from_str(&raw_mock_data)?;

            db.query(format!("REMOVE DATABASE {db_name}")).await?;
            for schema in DB_SCHEMA_FILES.map(std::fs::read_to_string) {
                db.query(schema?).await?;
            }
            surreal_brain::db::migration0(&db, &mock_data).await?;
            Ok::<(), anyhow::Error>(())
        })
        .await;
    Ok(db)
}

pub async fn run_service_in_background() -> Result<SocketAddr> {
    dotenv().ok();
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let addr = listener.local_addr()?;

    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?;
        let db_arc = Arc::new(db);

        Server::builder()
            .add_service(BrainGeneralCliServer::new(GeneralCliServer::new(db_arc.clone())))
            .add_service(BrainVmCliServer::new(VmCliServer::new(db_arc.clone())))
            .add_service(BrainVmDaemonServer::new(VmDaemonServer::new(db_arc.clone())))
            .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
            .await?;

        Ok::<(), anyhow::Error>(())
    });

    tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

    Ok(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?;
        let db_arc = Arc::new(db);

        tonic::transport::Server::builder()
            .add_service(BrainGeneralCliServer::new(GeneralCliServer::new(db_arc.clone())))
            .add_service(BrainVmCliServer::new(VmCliServer::new(db_arc.clone())))
            .add_service(BrainVmDaemonServer::new(VmDaemonServer::new(db_arc.clone())))
            .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
            .await?;

        Ok::<(), anyhow::Error>(())
    });
    client
}

pub async fn connect_stream_client_channel(c_stream: DuplexStream) -> Result<Channel> {
    let mut client = Some(c_stream);

    Ok(Endpoint::from_static("http://127.0.0.1:0")
        .connect_with_connector(service_fn(move |_: Uri| {
            let client = client.take().unwrap();
            async move { Ok::<TokioIo<DuplexStream>, std::io::Error>(TokioIo::new(client)) }
        }))
        .await?)
}

pub async fn run_service_for_stream() -> Result<Channel> {
    let client = run_service_for_stream_server().await;

    connect_stream_client_channel(client).await
}