brain/tests/grpc_test.rs

use common::prepare_test_env::{
    prepare_test_db, run_service_for_stream, run_service_in_background,
};
use common::test_utils::Key;
use common::vm_cli_utils::create_new_vm;
use common::vm_daemon_utils::mock_vm_daemon;
use detee_shared::common_proto::{Empty, Pubkey};
use detee_shared::general_proto::brain_general_cli_client::BrainGeneralCliClient;
use detee_shared::general_proto::{AirdropReq, ReportNodeReq};
use detee_shared::vm_proto::brain_vm_cli_client::BrainVmCliClient;
use detee_shared::vm_proto::ListVmContractsReq;
use futures::StreamExt;
use itertools::Itertools;
use std::vec;
use surreal_brain::constants::VM_NODE;
use surreal_brain::db::vm::VmNodeWithReports;

mod common;

#[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_general_airdrop() {
    // env_logger::builder().filter_level(log::LevelFilter::Trace).init();
    let _ = prepare_test_db().await;

    const AIRDROP_MULTIPLE: u64 = 1_000_000_000;
    let airdrop_amount = 10;

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

    let admin_keys = vec![Key::new(), Key::new(), Key::new()];
    let admin_pub_keys = admin_keys.iter().map(|k| k.pubkey.clone()).join(", ");
    std::env::set_var("ADMIN_PUB_KEYS", admin_pub_keys);

    let user_01_key = Key::new();
    let user_01_pubkey = user_01_key.pubkey.clone();

    let airdrop_req = AirdropReq { pubkey: user_01_pubkey.clone(), tokens: airdrop_amount };

    // user airdroping himself
    let err =
        client.airdrop(user_01_key.sign_request(airdrop_req.clone()).unwrap()).await.err().unwrap();
    assert_eq!(err.message(), "This operation is reserved to admin accounts");

    // other user airdroping
    let err =
        client.airdrop(Key::new().sign_request(airdrop_req.clone()).unwrap()).await.err().unwrap();
    assert_eq!(err.message(), "This operation is reserved to admin accounts");

    let _ = client.airdrop(admin_keys[0].sign_request(airdrop_req.clone()).unwrap()).await.unwrap();

    let bal_req_data = Pubkey { pubkey: user_01_pubkey };
    let bal_req = user_01_key.sign_request(bal_req_data.clone()).unwrap();
    let acc_bal_user_01 = client.get_balance(bal_req).await.unwrap().into_inner();

    assert_eq!(acc_bal_user_01.balance, airdrop_amount * AIRDROP_MULTIPLE);
    assert_eq!(acc_bal_user_01.tmp_locked, 0);

    // second airdrop from same admin
    let _ = client.airdrop(admin_keys[0].sign_request(airdrop_req.clone()).unwrap()).await.unwrap();

    let acc_bal_user_01 = client
        .get_balance(user_01_key.sign_request(bal_req_data.clone()).unwrap())
        .await
        .unwrap()
        .into_inner();

    assert_eq!(acc_bal_user_01.balance, 2 * airdrop_amount * AIRDROP_MULTIPLE);

    // third airdrop from another admin
    let _ = client.airdrop(admin_keys[1].sign_request(airdrop_req.clone()).unwrap()).await.unwrap();

    let acc_bal_user_01 = client
        .get_balance(user_01_key.sign_request(bal_req_data).unwrap())
        .await
        .unwrap()
        .into_inner();

    assert_eq!(acc_bal_user_01.balance, 3 * airdrop_amount * AIRDROP_MULTIPLE);

    // self airdrop
    let airdrop_req = AirdropReq { pubkey: admin_keys[2].pubkey.clone(), tokens: airdrop_amount };

    let _ = client.airdrop(admin_keys[2].sign_request(airdrop_req.clone()).unwrap()).await.unwrap();

    let bal_req_data = Pubkey { pubkey: admin_keys[2].pubkey.clone() };
    let acc_bal_admin_3 = client
        .get_balance(admin_keys[2].sign_request(bal_req_data.clone()).unwrap())
        .await
        .unwrap()
        .into_inner();

    assert_eq!(acc_bal_admin_3.balance, airdrop_amount * AIRDROP_MULTIPLE);
}

#[tokio::test]
async fn test_vm_creation() {
    let db = prepare_test_db().await;

    let brain_channel = run_service_for_stream().await;
    let daemon_key = mock_vm_daemon(brain_channel.clone()).await;

    let key = Key::new();

    let _ = create_new_vm(&db, key.clone(), daemon_key.clone(), brain_channel.clone()).await;
}

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

    let brain_channel = run_service_for_stream().await;
    let daemon_key = mock_vm_daemon(brain_channel.clone()).await;
    let mut client_gen_cli = BrainGeneralCliClient::new(brain_channel.clone());

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

    let report_req = ReportNodeReq {
        admin_pubkey: pubkey.clone(),
        node_pubkey: daemon_key.clone(),
        contract: String::from("uuid"),
        reason: String::from("reason"),
    };

    let report_error =
        client_gen_cli.report_node(key.sign_request(report_req).unwrap()).await.err().unwrap();

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

    let active_vm_id =
        create_new_vm(&db, key.clone(), daemon_key.clone(), brain_channel.clone()).await;

    let reason = String::from("something went wrong on vm");
    let report_req = ReportNodeReq {
        admin_pubkey: pubkey,
        node_pubkey: daemon_key.clone(),
        contract: active_vm_id,
        reason: reason.clone(),
    };

    let _ = client_gen_cli
        .report_node(key.sign_request(report_req).unwrap())
        .await
        .unwrap()
        .into_inner();

    let vm_nodes: Vec<VmNodeWithReports> = db
        .query(format!(
            "SELECT *, <-report.* as reports FROM {VM_NODE} WHERE id = {VM_NODE}:{daemon_key};"
        ))
        .await
        .unwrap()
        .take(0)
        .unwrap();

    let vm_node_with_report = vm_nodes.get(0).unwrap();

    assert!(vm_node_with_report.reports[0].reason == reason);
}

#[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
}