brain/tests/common/prepare_test_env.rs

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 surreal_brain::BrainState;
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<Client> {
    let db_connection = surreal_brain::db::db_connection(DB_URL, DB_NS, DB_NAME).await.unwrap();
    DB_STATE
        .get_or_init(|| async {
            // surreal_brain::db::init(DB_URL, DB_NS, DB_NAME)
            //     .await
            //     .expect("Failed to initialize the database");

            // surreal_brain::db::DB
            //     .set(surreal_brain::db::db_connection(DB_URL, DB_NS, DB_NAME).await.unwrap())
            //     .unwrap();

            let old_brain_data = surreal_brain::old_brain::BrainData::load_from_disk().unwrap();
            db_connection.query(format!("REMOVE DATABASE {DB_NAME}")).await.unwrap();
            db_connection
                .query(std::fs::read_to_string("interim_tables.surql").unwrap())
                .await
                .unwrap();
            surreal_brain::db::migration0(&db_connection, &old_brain_data).await.unwrap();
        })
        .await;
    db_connection
}

/*
use db::DB;

#[tokio::main]
async fn main() {
    env_logger::builder().filter_level(log::LevelFilter::Debug).init();
    // db::init(DB_ADDRESS, DB_NS, DB_NAME).await.unwrap();

    DB.set(db::db_connection(DB_ADDRESS, DB_NS, DB_NAME).await.unwrap()).unwrap();

    let db_connection = db::db_connection(DB_ADDRESS, DB_NS, DB_NAME).await.unwrap();
    let state = Arc::new(BrainState { db_connection });

    let addr = BRAIN_GRPC_ADDR.parse().unwrap();

    let snp_daemon_server = BrainVmDaemonServer::new(BrainVmDaemonForReal::new(state.clone()));
    let snp_cli_server = BrainVmCliServer::new(BrainVmCliForReal::new(state.clone()));
    let general_service_server =
        BrainGeneralCliServer::new(BrainGeneralCliForReal::new(state.clone()));

    let cert = std::fs::read_to_string(CERT_PATH).unwrap();
    let key = std::fs::read_to_string(CERT_KEY_PATH).unwrap();

    let identity = Identity::from_pem(cert, key);

    Server::builder()
        .tls_config(ServerTlsConfig::new().identity(identity))
        .unwrap()
        .add_service(snp_daemon_server)
        .add_service(snp_cli_server)
 */

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_connection = surreal_brain::db::db_connection(DB_URL, DB_NS, DB_NAME).await.unwrap();
        let state = Arc::new(BrainState { db_connection });

        Server::builder()
            .add_service(BrainGeneralCliServer::new(BrainGeneralCliForReal::new(state.clone())))
            .add_service(BrainVmCliServer::new(BrainVmCliForReal::new(state.clone())))
            .add_service(BrainVmDaemonServer::new(BrainVmDaemonForReal::new(state.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_connection = surreal_brain::db::db_connection(DB_URL, DB_NS, DB_NAME).await.unwrap();
        let state = Arc::new(BrainState { db_connection });

        tonic::transport::Server::builder()
            .add_service(BrainGeneralCliServer::new(BrainGeneralCliForReal::new(state.clone())))
            .add_service(BrainVmCliServer::new(BrainVmCliForReal::new(state.clone())))
            .add_service(BrainVmDaemonServer::new(BrainVmDaemonForReal::new(state.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::<TokioIo<DuplexStream>, 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
}