brain/tests/grpcs_test.rs

use detee_shared::common_proto::Empty;
use detee_shared::general_proto::ReportNodeReq;
use detee_shared::vm_proto::brain_vm_cli_client::BrainVmCliClient;
use detee_shared::vm_proto::brain_vm_daemon_client::BrainVmDaemonClient;
use detee_shared::vm_proto::{ListVmContractsReq, RegisterVmNodeReq};
use detee_shared::{
    common_proto::Pubkey, general_proto::brain_general_cli_client::BrainGeneralCliClient, vm_proto,
};
mod common;
use common::prepare_test_env::{
    connect_stream_client_channel, prepare_test_db, run_service_for_stream,
    run_service_for_stream_server, run_service_in_background,
};
use common::test_utils::Key;
use futures::StreamExt;
use tokio::task::JoinSet;
use tokio_stream::wrappers::ReceiverStream;

#[tokio::test]
async fn test_reg_vm_node() {
    prepare_test_db().await;

    let addr = run_service_in_background().await;
    let mut client = BrainVmDaemonClient::connect(format!("http://{}", addr)).await.unwrap();

    let operator_wallet = Key::new().pubkey;

    let key = Key::new();

    let node_pubkey = key.pubkey.clone();

    let req = RegisterVmNodeReq {
        node_pubkey,
        operator_wallet,
        main_ip: String::from("185.243.218.213"),
        city: String::from("Oslo"),
        country: String::from("Norway"),
        region: String::from("EU"),
        price: 1200,
    };

    let mut grpc_stream =
        client.register_vm_node(key.sign_request(req).unwrap()).await.unwrap().into_inner();

    let mut vm_contracts = Vec::new();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(vm_c) => {
                vm_contracts.push(vm_c);
            }
            Err(e) => {
                panic!("Received error instead of vm_contracts: {e:?}");
            }
        }
    }

    assert!(vm_contracts.is_empty())
}

#[tokio::test]
async fn test_brain_message() {
    // spawn grpc stream server
    // mock a daemon
    // mock a cli client to interact with brain

    // validate if something happening in "surreal_brain::db::NewVmReq", "surreal_brain::db::UpdateVmReq", "surreal_brain::db::DeletedVm" these table
    // mock daemon will responde to brain

    // TODO: register node before connecting to brain

    env_logger::builder().filter_level(log::LevelFilter::Info).init();
    let _ = prepare_test_db().await;

    let (tokio_duplex, addr) = run_service_for_stream_server().await;

    let channel = connect_stream_client_channel(tokio_duplex, addr).await;

    let mut daemon_client = BrainVmDaemonClient::new(channel.clone());

    let daemon_key = Key::new();

    let mut daemon_join_set = JoinSet::new();

    let (tx, mut brain_msg_rx) = tokio::sync::mpsc::channel(1);

    // listen to brain
    let mut daemon_client_01 = daemon_client.clone();
    let daemon_key_01 = daemon_key.clone();
    daemon_join_set.spawn(async move {
        let mut grpc_stream = daemon_client_01
            .brain_messages(daemon_key_01.sign_stream_auth(vec![]).unwrap())
            .await
            .unwrap()
            .into_inner();

        while let Some(Ok(stream_update)) = grpc_stream.next().await {
            log::info!("vm deamon got notified: {:?}", &stream_update);
            let _ = tx.send(stream_update).await;
        }
    });

    // send to brain
    let (daemon_msg_tx, rx) = tokio::sync::mpsc::channel(1);
    let daemon_msg_tx_01 = daemon_msg_tx.clone();
    let daemon_key_02 = daemon_key.clone();
    daemon_join_set.spawn(async move {
        let rx_stream = ReceiverStream::new(rx);
        daemon_msg_tx_01
            .send(vm_proto::VmDaemonMessage {
                msg: Some(vm_proto::vm_daemon_message::Msg::Auth(
                    daemon_key_02.sign_stream_auth(vec![]).unwrap(),
                )),
            })
            .await
            .unwrap();
        daemon_client.daemon_messages(rx_stream).await.unwrap();
    });

    // daemon engine
    daemon_join_set.spawn(async move {
        while let Some(brain_msg) = brain_msg_rx.recv().await {
            match brain_msg.msg {
                Some(vm_proto::brain_vm_message::Msg::NewVmReq(new_vm_req)) => {
                    let args = Some(vm_proto::MeasurementArgs {
                        dtrfs_api_endpoint: String::from("184.107.169.199:48865"),
                        exposed_ports: new_vm_req.extra_ports,
                        ovmf_hash: String::from(
                            "0346619257269b9a61ee003e197d521b8e2283483070d163a34940d6a1d40d76",
                        ),
                        ips: vec![],
                    });

                    let new_vm_resp = vm_proto::NewVmResp {
                        uuid: new_vm_req.uuid.clone(),
                        args,
                        error: String::new(),
                    };

                    let res_data = vm_proto::VmDaemonMessage {
                        msg: Some(vm_proto::vm_daemon_message::Msg::NewVmResp(new_vm_resp)),
                    };
                    daemon_msg_tx.send(res_data).await.unwrap();
                }
                Some(vm_proto::brain_vm_message::Msg::UpdateVmReq(_update_vm_req)) => {
                    //
                }
                Some(vm_proto::brain_vm_message::Msg::DeleteVm(_del_vm_req)) => {
                    //
                }
                None => todo!(),
            }
        }
    });

    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;

    let mut cli_client = BrainVmCliClient::new(channel);

    let cli_key = Key::new();

    let req = vm_proto::NewVmReq {
        admin_pubkey: cli_key.pubkey.clone(),
        node_pubkey: daemon_key.pubkey.clone(),
        price_per_unit: 1200,
        extra_ports: vec![8080, 8081],
        locked_nano: 0,
        ..Default::default()
    };
    let new_vm_resp =
        cli_client.new_vm(cli_key.sign_request(req).unwrap()).await.unwrap().into_inner();

    assert!(new_vm_resp.error.is_empty());
    assert!(new_vm_resp.uuid.len() == 40);

    let id = ("measurement_args", new_vm_resp.uuid);
    let data_in_db: detee_shared::vm_proto::MeasurementArgs =
        surreal_brain::db::DB.select(id).await.unwrap().unwrap();

    assert_eq!(data_in_db, new_vm_resp.args.unwrap());
}

#[tokio::test]
async fn test_general_balance() {
    // env_logger::builder().filter_level(log::LevelFilter::Trace).init();
    let _ = prepare_test_db().await;

    let addr = run_service_in_background().await;
    let mut client = BrainGeneralCliClient::connect(format!("http://{}", addr)).await.unwrap();

    let key = Key::new();
    let pubkey = key.pubkey.clone();
    let req_data = Pubkey { pubkey };

    let req = key.sign_request(req_data).unwrap();

    let acc_bal = client.get_balance(req).await.unwrap().into_inner();

    assert_eq!(acc_bal.balance, 0);
    assert_eq!(acc_bal.tmp_locked, 0);
}

#[tokio::test]
async fn test_report_node() {
    prepare_test_db().await;

    let addr = run_service_in_background().await;
    let mut client = BrainGeneralCliClient::connect(format!("http://{}", addr)).await.unwrap();

    let key = Key::new();
    let pubkey = key.pubkey.clone();

    // TODO: create contract, node and operator in db and use it here
    let req_data = ReportNodeReq {
        admin_pubkey: pubkey,
        node_pubkey: String::from("node_pubkey"),
        contract: String::from("uuid"),
        reason: String::from("reason"),
    };

    let report_error = client.report_node(key.sign_request(req_data).unwrap()).await.err().unwrap();

    println!("Report error: {:?}", report_error);
    assert_eq!(report_error.message(), "No contract found by this ID.");

    // verify report in db
}

#[tokio::test]
// TODO: register some operators before testing this
async fn test_list_operators() {
    prepare_test_db().await;

    let channel = run_service_for_stream().await;

    let mut client = BrainGeneralCliClient::new(channel);

    let key = Key::new();

    let mut grpc_stream =
        client.list_operators(key.sign_request(Empty {}).unwrap()).await.unwrap().into_inner();

    let mut operators = Vec::new();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(op) => {
                operators.push(op);
            }
            Err(e) => {
                panic!("Received error instead of operators: {e:?}");
            }
        }
    }

    assert!(!operators.is_empty())
}

#[tokio::test]
// TODO: create vm for this user before testing this
async fn test_list_vm_contracts() {
    prepare_test_db().await;

    let channel = run_service_for_stream().await;
    let mut client = BrainVmCliClient::new(channel);

    let key = Key::new();
    let pubkey = key.pubkey.clone();

    let req_data =
        ListVmContractsReq { wallet: pubkey, uuid: String::from("uuid"), as_operator: false };

    let mut grpc_stream =
        client.list_vm_contracts(key.sign_request(req_data).unwrap()).await.unwrap().into_inner();

    let mut vm_contracts = Vec::new();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(vm_c) => {
                vm_contracts.push(vm_c);
            }
            Err(e) => {
                panic!("Received error instead of vm_contracts: {e:?}");
            }
        }
    }

    assert!(vm_contracts.is_empty())

    // verify report in db
}