brain/tests/grpc_vm_cli_test.rs
Noor 9363750d5b
implemented redirect vm methods
Introduces new environment variable to set public endpoint of the brain
refactor redirect validation to db module
validating pubsub node on all vm
fixed some vm tests
2025-06-12 18:44:45 +05:30

303 lines
11 KiB
Rust

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;
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::{DeleteVmReq, ExtendVmReq, ListVmContractsReq, NewVmReq};
use futures::StreamExt;
use surreal_brain::constants::{ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, TOKEN_DECIMAL};
use surreal_brain::db::prelude as db;
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..