brain/tests/grpc_vm_cli_test.rs

use common::prepare_test_env::{prepare_test_db, run_service_for_stream};
use common::test_utils::{airdrop, Key};
use common::vm_cli_utils::{create_new_vm, user_list_vm_contracts};
use common::vm_daemon_utils::{mock_vm_daemon, register_vm_node};
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::{self, DeleteVmReq, RegisterVmNodeReq};
use detee_shared::vm_proto::{ExtendVmReq, ListVmContractsReq, NewVmReq};
use futures::StreamExt;
use std::vec;
use surreal_brain::constants::{
    ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, TOKEN_DECIMAL, VM_NODE,
};
use surreal_brain::db::prelude as db;
use surreal_brain::grpc::check_pubsub_node;

mod common;

#[tokio::test]
async fn test_vm_creation() {
    let db = prepare_test_db().await.unwrap();
    /*
    env_logger::builder()
        .filter_level(log::LevelFilter::Debug)
        .filter_module("tungstenite", log::LevelFilter::Debug)
        .filter_module("tokio_tungstenite", log::LevelFilter::Debug)
        .init();
    */

    let brain_channel = run_service_for_stream().await.unwrap();
    let daemon_key = mock_vm_daemon(&brain_channel, None).await.unwrap();

    let key = Key::new();

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

    assert!(new_vm_resp.is_err());
    let grpc_error_message = new_vm_resp.err().unwrap().to_string();
    assert!(grpc_error_message.contains("Insufficient funds"));

    let airdrop_amount = 10;
    airdrop(&brain_channel, &key.pubkey, airdrop_amount).await.unwrap();

    let new_vm_id = create_new_vm(&db, &key, &daemon_key, &brain_channel).await.unwrap();
    let active_vm: Option<db::ActiveVm> = db.select((ACTIVE_VM, new_vm_id)).await.unwrap();
    assert!(active_vm.is_some());

    let daemon_key_02 =
        mock_vm_daemon(&brain_channel, Some("something went wrong 04".to_string())).await.unwrap();

    let locking_nano = 100;
    let mut new_vm_req = vm_proto::NewVmReq {
        admin_pubkey: key.pubkey.clone(),
        node_pubkey: daemon_key_02.clone(),
        price_per_unit: 1200,
        extra_ports: vec![8080, 8081],
        locked_nano: locking_nano,
        ..Default::default()
    };

    let mut client_vm_cli = BrainVmCliClient::new(brain_channel.clone());
    let new_vm_resp = client_vm_cli
        .new_vm(key.sign_request(new_vm_req.clone()).unwrap())
        .await
        .unwrap()
        .into_inner();

    assert!(!new_vm_resp.error.is_empty());

    let vm_req_db: Option<db::NewVmReq> =
        db.select((NEW_VM_REQ, new_vm_resp.uuid.clone())).await.unwrap();

    if let Some(new_vm_req) = vm_req_db {
        assert!(!new_vm_req.error.is_empty());
        assert_eq!(new_vm_resp.error, new_vm_req.error);
    }

    let acc_db: db::Account = db.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
    assert_eq!(acc_db.balance, (airdrop_amount * TOKEN_DECIMAL - 100));
    assert_eq!(acc_db.tmp_locked, 0);

    new_vm_req.node_pubkey = daemon_key;

    let new_vm_resp =
        client_vm_cli.new_vm(key.sign_request(new_vm_req).unwrap()).await.unwrap().into_inner();
    assert!(new_vm_resp.error.is_empty());

    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
    let active_vm: Option<db::ActiveVm> = db.select((ACTIVE_VM, new_vm_resp.uuid)).await.unwrap();
    assert!(active_vm.is_some());

    let acc_db: db::Account = db.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
    assert_eq!(acc_db.balance, airdrop_amount * TOKEN_DECIMAL - (locking_nano + 100));
    assert_eq!(acc_db.tmp_locked, 0);
}

#[tokio::test]
async fn test_timeout_vm_creation() {
    prepare_test_db().await.unwrap();
    // env_logger::builder().filter_level(log::LevelFilter::Error).init();

    let brain_channel = run_service_for_stream().await.unwrap();
    let mut daemon_client = BrainVmDaemonClient::new(brain_channel.clone());
    let daemon_key = Key::new();

    register_vm_node(&mut daemon_client, &daemon_key, &Key::new().pubkey).await.unwrap();

    let key = Key::new();

    let new_vm_req = NewVmReq {
        admin_pubkey: key.pubkey.clone(),
        node_pubkey: daemon_key.pubkey,
        price_per_unit: 1200,
        extra_ports: vec![8080, 8081],
        locked_nano: 100,
        ..Default::default()
    };

    airdrop(&brain_channel, &key.pubkey, 10).await.unwrap();

    let mut client_vm_cli = BrainVmCliClient::new(brain_channel.clone());
    let timeout_error =
        client_vm_cli.new_vm(key.sign_request(new_vm_req).unwrap()).await.err().unwrap();

    assert_eq!(
        timeout_error.message(),
        "Network timeout. Please try again later or contact the DeTEE devs team.",
    )
}

#[tokio::test]
async fn test_vm_deletion() {
    let db = prepare_test_db().await.unwrap();
    let brain_channel = run_service_for_stream().await.unwrap();
    let daemon_key = mock_vm_daemon(&brain_channel, None).await.unwrap();

    let key = Key::new();
    airdrop(&brain_channel, &key.pubkey, 10).await.unwrap();

    let new_vm_uuid = create_new_vm(&db, &key, &daemon_key, &brain_channel).await.unwrap();

    let mut client_vm_cli = BrainVmCliClient::new(brain_channel.clone());

    let del_app_req = DeleteVmReq { admin_pubkey: key.pubkey.clone(), uuid: new_vm_uuid };
    let _ = client_vm_cli.delete_vm(key.sign_request(del_app_req).unwrap()).await.unwrap();

    let key_02 = Key::new();

    // delete random vm
    let mut del_vm_req = DeleteVmReq {
        admin_pubkey: key_02.pubkey.clone(),
        uuid: "9ae3VH8nJg2i8pqTQ6mJtvYuS2kd9n1XLLco8GUPfT95".to_string(),
    };

    let del_err = client_vm_cli
        .delete_vm(key_02.sign_request(del_vm_req.clone()).unwrap())
        .await
        .err()
        .unwrap();
    assert_eq!(del_err.message(), "Unauthorized");

    let new_vm_uuid_02 = create_new_vm(&db, &key, &daemon_key, &brain_channel).await.unwrap();
    del_vm_req.uuid = new_vm_uuid_02;

    let del_err = client_vm_cli
        .delete_vm(key_02.sign_request(del_vm_req.clone()).unwrap())
        .await
        .err()
        .unwrap();
    assert_eq!(del_err.message(), "Unauthorized");

    // test refund
    let key_03 = Key::new();
    airdrop(&brain_channel, &key_03.pubkey, 10).await.unwrap();

    let new_vm_uuid_03 = create_new_vm(&db, &key_03, &daemon_key, &brain_channel).await.unwrap();

    let del_vm_req =
        DeleteVmReq { admin_pubkey: key_03.pubkey.clone(), uuid: new_vm_uuid_03.clone() };
    let _ = client_vm_cli.delete_vm(key_03.sign_request(del_vm_req).unwrap()).await.unwrap();

    let acc: db::Account = db.select((ACCOUNT, key_03.pubkey.clone())).await.unwrap().unwrap();
    assert_eq!(acc.balance, 10 * TOKEN_DECIMAL);

    let deleted_vm =
        db.select::<Option<db::DeletedVm>>((DELETED_VM, new_vm_uuid_03)).await.unwrap();
    assert!(deleted_vm.is_some());
}

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

    let channel = run_service_for_stream().await.unwrap();
    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
}

/*
Private Key: 69kPgwAzftzNf5oBFQwqocFAgsn17ZwKxCHFvpoZ1LRn
Public Key : HYo88ncAPekykFDTFGdJV9ehJ81LL4onQ8xzYRAu2Ph4

Public Key : Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG // operator
Private Key: 2oQvpUMSzaZmaVLfVWsriRbhRK1JbnPEBed5UmpZSDNc */

#[tokio::test]
async fn test_extend_vm() {
    let db_conn = prepare_test_db().await.unwrap();

    let channel = run_service_for_stream().await.unwrap();
    let mut cli_client = BrainVmCliClient::new(channel.clone());

    /*
    env_logger::builder()
        .filter_level(log::LevelFilter::Debug)
        .filter_module("tungstenite", log::LevelFilter::Debug)
        .filter_module("tokio_tungstenite", log::LevelFilter::Debug)
        .init();
     */

    let key = Key::from("YnaZccN852KYvnhV5pbeBrHV4dSwgzpQXghdyMiHC6i"); // GmE4JH3bL4NpmzwKCBJemJzRTumJAnbcXLGqce5mREgS admin
    let vm_contract_id = "1d749a816c27a22efa0e574b023a6afef040fe5";

    let locking_01 = 200;
    let locking_02 = 86000000000u64;

    let mut req = ExtendVmReq {
        admin_pubkey: key.pubkey.clone(),
        uuid: "foooooooo".to_string(),
        locked_nano: u64::MAX,
    };

    let err_precondition =
        cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.err().unwrap();
    assert!(err_precondition.to_string().contains("Transation Too Big "));

    req.locked_nano = 99999999999999;
    let err_precondition =
        cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.err().unwrap();
    assert!(err_precondition.to_string().contains("Contract not found"));

    req.uuid = vm_contract_id.to_string();
    let err_precondition =
        cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.err().unwrap();
    assert!(err_precondition.to_string().contains("Insufficient funds"));

    req.locked_nano = locking_01;
    let err_precondition =
        cli_client.extend_vm(Key::new().sign_request(req.clone()).unwrap()).await.err().unwrap();
    assert!(err_precondition.to_string().contains("signature does not match"));

    let acc: db::Account = db_conn.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
    let contract: db::ActiveVm =
        db_conn.select((ACTIVE_VM, vm_contract_id)).await.unwrap().unwrap();

    let expected_bal_01 =
        (acc.balance as i128 + (contract.locked_nano as i128 - locking_01 as i128)) as u64;

    cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.unwrap();

    let contract = &user_list_vm_contracts(&channel, &key, false, vm_contract_id).await.unwrap()[0];
    assert_eq!(contract.locked_nano, locking_01);

    let acc: db::Account = db_conn.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
    assert_eq!(acc.balance, expected_bal_01);

    let expected_bal_02 =
        (acc.balance as i128 + (contract.locked_nano as i128 - locking_02 as i128)) as u64;

    req.locked_nano = locking_02;

    cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.unwrap();

    let contract = &user_list_vm_contracts(&channel, &key, false, vm_contract_id).await.unwrap()[0];
    assert_eq!(contract.locked_nano, locking_02);

    let acc: db::Account = db_conn.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
    assert_eq!(acc.balance, expected_bal_02);
}

// TODO: test register vm node, delete vm contract while node offline, kick, etc..

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

    let brain_channel = run_service_for_stream().await.unwrap();
    let mut vm_daemon_client = BrainVmDaemonClient::new(brain_channel.clone());

    let node_key = Key::new();
    let operator_wallet = "Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG";

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

    vm_daemon_client
        .register_vm_node(node_key.sign_request(req).unwrap())
        .await
        .unwrap()
        .into_inner();

    // let client_key = Key::new();

    // let new_vm_req = vm_proto::NewVmReq {
    //     admin_pubkey: client_key.pubkey.clone(),
    //     node_pubkey: node_key.pubkey.clone(),
    //     price_per_unit: 1200,
    //     extra_ports: vec![8080, 8081],
    //     locked_nano: 100,
    //     ..Default::default()
    // };

    let id = surrealdb::RecordId::from((VM_NODE, node_key.pubkey.clone()));

    let endpoint = check_pubsub_node(&db, id).await;
    dbg!(&endpoint);

    // let mut client_vm_cli = BrainVmCliClient::new(brain_channel.clone());
    // let redirect = client_vm_cli.new_vm(client_key.sign_request(new_vm_req).unwrap()).await;
}