brain/tests/common/vm_cli_utils.rs
Noor af18e4ee77
fixes and tests
admin keys from env
test operator inspection
test vm creation timeout
extensive tests on airdrop
refactor db module imports
fix register vm_node creates operator account in db
modularised test into its module and reusable method
fix test brain message add ssh port on mock daemon while new vm
improved error handling on tests unwraping only on top level method
test utils methods accepts refs to remove clone() on top level methods
2025-05-08 15:12:58 +05:30

81 lines
2.6 KiB
Rust

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::{AirdropReq, ReportNodeReq};
use detee_shared::vm_proto;
use detee_shared::vm_proto::brain_vm_cli_client::BrainVmCliClient;
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;
async fn airdrop(brain_channel: &Channel, wallet: &str, amount: u64) -> Result<()> {
let mut client = BrainGeneralCliClient::new(brain_channel.clone());
let airdrop_req = AirdropReq { pubkey: wallet.to_string(), tokens: amount };
let admin_key = Key::new();
std::env::set_var("ADMIN_PUB_KEYS", &admin_key.pubkey);
client.airdrop(admin_key.sign_request(airdrop_req.clone())?).await?;
Ok(())
}
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: 0,
..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.uuid.len() == 40);
// wait for update db
tokio::time::sleep(tokio::time::Duration::from_millis(700)).await;
let vm_req_db: Option<db::NewVmReq> = db.select((NEW_VM_REQ, new_vm_resp.uuid.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.uuid.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?)
}