brain/tests/common/vm_cli_utils.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-License-Identifier: Unlicense

use super::test_utils::Key;
use anyhow::{anyhow, Result};
use detee_shared::common_proto::Empty;
use detee_shared::general_proto::brain_general_cli_client::BrainGeneralCliClient;
use detee_shared::general_proto::{Account, RegOperatorReq, ReportNodeReq};
use detee_shared::vm_proto::brain_vm_cli_client::BrainVmCliClient;
use detee_shared::{app_proto, vm_proto};
use futures::StreamExt;
use surreal_brain::constants::{ACTIVE_VM, NEW_VM_REQ};
use surreal_brain::db::prelude as db;
use surrealdb::engine::remote::ws::Client;
use surrealdb::Surreal;
use tonic::transport::Channel;

pub async fn create_new_vm(
    db: &Surreal<Client>,
    key: &Key,
    node_pubkey: &str,
    brain_channel: &Channel,
) -> Result<String> {
    let new_vm_req = vm_proto::NewVmReq {
        admin_pubkey: key.pubkey.clone(),
        node_pubkey: node_pubkey.to_string(),
        price_per_unit: 1200,
        extra_ports: vec![8080, 8081],
        locked_nano: 100,
        ..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)?).await?.into_inner();

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

    // wait for update db
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

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

    if let Some(new_vm_req) = vm_req_db {
        panic!("New VM request found in DB: {:?}", new_vm_req);
    }

    let active_vm_op: Option<db::ActiveVm> =
        db.select((ACTIVE_VM, new_vm_resp.vm_id.clone())).await?;
    let active_vm = active_vm_op.ok_or(anyhow!("Not found active vm in db"))?;

    Ok(active_vm.id.key().to_string())
}

pub async fn report_node(
    key: &Key,
    brain_channel: &Channel,
    node_pubkey: &str,
    contract: &str,
    reason: &str,
) -> Result<tonic::Response<Empty>> {
    let mut client_gen_cli = BrainGeneralCliClient::new(brain_channel.clone());

    let report_req = ReportNodeReq {
        admin_pubkey: key.pubkey.clone(),
        node_pubkey: node_pubkey.to_string(),
        contract: contract.to_string(),
        reason: reason.to_string(),
    };

    Ok(client_gen_cli.report_node(key.sign_request(report_req)?).await?)
}

pub async fn register_operator(brain_channel: &Channel, key: &Key, escrow_nano: u64) -> Result<()> {
    let mut cli_client = BrainGeneralCliClient::new(brain_channel.clone());
    let reg_req = RegOperatorReq {
        pubkey: key.pubkey.clone(),
        escrow: escrow_nano,
        email: "foo@bar.com".to_string(),
    };

    cli_client.register_operator(key.sign_request(reg_req.clone()).unwrap()).await?;
    Ok(())
}

pub async fn list_accounts(brain_channel: &Channel, admin_key: &Key) -> Result<Vec<Account>> {
    let mut cli_client = BrainGeneralCliClient::new(brain_channel.clone());
    let mut stream =
        cli_client.list_accounts(admin_key.sign_request(Empty {}).unwrap()).await?.into_inner();

    let mut accounts = Vec::new();

    while let Some(stream_data) = stream.next().await {
        match stream_data {
            Ok(account) => {
                accounts.push(account);
            }
            Err(e) => {
                panic!("Error while listing accounts: {e:?}");
            }
        }
    }

    Ok(accounts)
}

pub async fn list_all_vm_contracts(
    brain_channel: &Channel,
    admin_key: &Key,
) -> Result<Vec<vm_proto::VmContract>> {
    let mut cli_client = BrainGeneralCliClient::new(brain_channel.clone());
    let mut stream = cli_client
        .list_all_vm_contracts(admin_key.sign_request(Empty {}).unwrap())
        .await?
        .into_inner();

    let mut vm_contracts = Vec::new();

    while let Some(stream_data) = stream.next().await {
        match stream_data {
            Ok(vm_contract) => {
                vm_contracts.push(vm_contract);
            }
            Err(e) => {
                panic!("Error while listing vm_contracts: {e:?}");
            }
        }
    }

    Ok(vm_contracts)
}

pub async fn list_all_app_contracts(
    brain_channel: &Channel,
    admin_key: &Key,
) -> Result<Vec<app_proto::AppContract>> {
    let mut cli_client = BrainGeneralCliClient::new(brain_channel.clone());
    let mut stream = cli_client
        .list_all_app_contracts(admin_key.sign_request(Empty {}).unwrap())
        .await?
        .into_inner();

    let mut app_contracts = Vec::new();

    while let Some(stream_data) = stream.next().await {
        match stream_data {
            Ok(app_contract) => {
                app_contracts.push(app_contract);
            }
            Err(e) => {
                panic!("Error while listing app_contracts: {e:?}");
            }
        }
    }

    Ok(app_contracts)
}

pub async fn user_list_vm_contracts(
    brain_channel: &Channel,
    key: &Key,
    as_operator: bool,
    vm_id: &str,
) -> Result<Vec<vm_proto::VmContract>> {
    let mut cli_client = BrainVmCliClient::new(brain_channel.clone());
    let mut stream = cli_client
        .list_vm_contracts(
            key.sign_request(vm_proto::ListVmContractsReq {
                wallet: key.pubkey.clone(),
                as_operator,
                vm_id: vm_id.to_string(),
            })
            .unwrap(),
        )
        .await?
        .into_inner();

    let mut vm_contracts = Vec::new();

    while let Some(stream_data) = stream.next().await {
        match stream_data {
            Ok(vm_contract) => {
                vm_contracts.push(vm_contract);
            }
            Err(e) => {
                panic!("Error while listing vm_contracts: {e:?}");
            }
        }
    }

    Ok(vm_contracts)
}