diff --git a/tests/common/test_utils.rs b/tests/common/test_utils.rs index af25ecd..5ae3fae 100644 --- a/tests/common/test_utils.rs +++ b/tests/common/test_utils.rs @@ -1,9 +1,24 @@ use anyhow::Result; use detee_shared::vm_proto as snp_proto; use ed25519_dalek::{Signer, SigningKey}; +use itertools::Itertools; +use std::sync::OnceLock; use tonic::metadata::AsciiMetadataValue; use tonic::Request; +pub static ADMIN_KEYS: OnceLock> = OnceLock::new(); + +pub fn admin_keys() -> Vec { + let admin_keys = ADMIN_KEYS.get_or_init(|| { + 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); + admin_keys.clone() + }); + + admin_keys.clone() +} + #[derive(Debug, Clone)] pub struct Key { pub sg_key: SigningKey, diff --git a/tests/common/vm_cli_utils.rs b/tests/common/vm_cli_utils.rs index 2d7c2c3..3a1bb40 100644 --- a/tests/common/vm_cli_utils.rs +++ b/tests/common/vm_cli_utils.rs @@ -1,4 +1,4 @@ -use super::test_utils::Key; +use super::test_utils::{admin_keys, Key}; use anyhow::{anyhow, Result}; use detee_shared::common_proto::Empty; use detee_shared::general_proto::brain_general_cli_client::BrainGeneralCliClient; @@ -11,12 +11,11 @@ use surrealdb::engine::remote::ws::Client; use surrealdb::Surreal; use tonic::transport::Channel; -async fn airdrop(brain_channel: &Channel, wallet: &str, amount: u64) -> Result<()> { +pub 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); + let admin_key = admin_keys()[0].clone(); client.airdrop(admin_key.sign_request(airdrop_req.clone())?).await?; diff --git a/tests/grpc_general_test.rs b/tests/grpc_general_test.rs index 6fa4ad6..2c601a0 100644 --- a/tests/grpc_general_test.rs +++ b/tests/grpc_general_test.rs @@ -1,17 +1,15 @@ 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, report_node}; +use common::test_utils::{admin_keys, Key}; +use common::vm_cli_utils::{airdrop, create_new_vm, report_node}; use common::vm_daemon_utils::{mock_vm_daemon, register_vm_node}; use detee_shared::common_proto::{Empty, Pubkey}; use detee_shared::general_proto::brain_general_cli_client::BrainGeneralCliClient; use detee_shared::general_proto::AirdropReq; use detee_shared::vm_proto::brain_vm_daemon_client::BrainVmDaemonClient; use futures::StreamExt; -use itertools::Itertools; -use std::vec; -use surreal_brain::constants::VM_NODE; +use surreal_brain::constants::{TOKEN_DECIMAL, VM_NODE}; use surreal_brain::db::vm::VmNodeWithReports; mod common; @@ -41,15 +39,12 @@ async fn test_general_airdrop() { // env_logger::builder().filter_level(log::LevelFilter::Trace).init(); prepare_test_db().await.unwrap(); - const AIRDROP_MULTIPLE: u64 = 1_000_000_000; let airdrop_amount = 10; let addr = run_service_in_background().await.unwrap(); 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 admin_keys = admin_keys(); let user_01_key = Key::new(); let user_01_pubkey = user_01_key.pubkey.clone(); @@ -72,7 +67,7 @@ async fn test_general_airdrop() { 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.balance, airdrop_amount * TOKEN_DECIMAL); assert_eq!(acc_bal_user_01.tmp_locked, 0); // second airdrop from same admin @@ -84,7 +79,7 @@ async fn test_general_airdrop() { .unwrap() .into_inner(); - assert_eq!(acc_bal_user_01.balance, 2 * airdrop_amount * AIRDROP_MULTIPLE); + assert_eq!(acc_bal_user_01.balance, 2 * airdrop_amount * TOKEN_DECIMAL); // third airdrop from another admin let _ = client.airdrop(admin_keys[1].sign_request(airdrop_req.clone()).unwrap()).await.unwrap(); @@ -95,7 +90,7 @@ async fn test_general_airdrop() { .unwrap() .into_inner(); - assert_eq!(acc_bal_user_01.balance, 3 * airdrop_amount * AIRDROP_MULTIPLE); + assert_eq!(acc_bal_user_01.balance, 3 * airdrop_amount * TOKEN_DECIMAL); // self airdrop let airdrop_req = AirdropReq { pubkey: admin_keys[2].pubkey.clone(), tokens: airdrop_amount }; @@ -109,7 +104,7 @@ async fn test_general_airdrop() { .unwrap() .into_inner(); - assert_eq!(acc_bal_admin_3.balance, airdrop_amount * AIRDROP_MULTIPLE); + assert_eq!(acc_bal_admin_3.balance, airdrop_amount * TOKEN_DECIMAL); } #[tokio::test] @@ -127,6 +122,8 @@ async fn test_report_node() { log::info!("Report error: {:?}", report_error); assert!(report_error.to_string().contains("No contract found by this ID.")); + airdrop(&brain_channel, &key.pubkey, 10).await.unwrap(); + let active_vm_id = create_new_vm(&db, &key, &daemon_key, &brain_channel).await.unwrap(); let reason = String::from("something went wrong on vm"); diff --git a/tests/grpc_vm_cli_test.rs b/tests/grpc_vm_cli_test.rs index 4b6069d..b6d4c71 100644 --- a/tests/grpc_vm_cli_test.rs +++ b/tests/grpc_vm_cli_test.rs @@ -1,12 +1,14 @@ use common::prepare_test_env::{prepare_test_db, run_service_for_stream}; use common::test_utils::Key; -use common::vm_cli_utils::create_new_vm; +use common::vm_cli_utils::{airdrop, create_new_vm}; 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::{ListVmContractsReq, NewVmReq}; use futures::StreamExt; use std::vec; +use surreal_brain::constants::ACTIVE_VM; +use surreal_brain::db::vm::ActiveVm; mod common; @@ -26,7 +28,11 @@ async fn test_vm_creation() { let grpc_error_message = new_vm_resp.err().unwrap().to_string(); assert!(grpc_error_message.contains("Insufficient funds")); - // TODO: Airdrop the user and try creating the VM again + airdrop(&brain_channel, &key.pubkey, 10).await.unwrap(); + + let new_vm_id = create_new_vm(&db, &key, &daemon_key, &brain_channel).await.unwrap(); + let active_vm: Option = db.select((ACTIVE_VM, new_vm_id)).await.unwrap(); + assert!(active_vm.is_some()); } #[tokio::test] @@ -51,6 +57,8 @@ async fn test_timeout_vm_creation() { ..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(); diff --git a/tests/grpc_vm_daemon_test.rs b/tests/grpc_vm_daemon_test.rs index 89d7983..c68c5da 100644 --- a/tests/grpc_vm_daemon_test.rs +++ b/tests/grpc_vm_daemon_test.rs @@ -2,6 +2,7 @@ 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::airdrop; 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; @@ -29,7 +30,7 @@ async fn test_brain_message() { let brain_channel = run_service_for_stream().await.unwrap(); let daemon_key = mock_vm_daemon(&brain_channel).await.unwrap(); - let mut cli_client = BrainVmCliClient::new(brain_channel); + let mut cli_client = BrainVmCliClient::new(brain_channel.clone()); let cli_key = Key::new(); @@ -41,6 +42,7 @@ async fn test_brain_message() { locked_nano: 0, ..Default::default() }; + airdrop(&brain_channel, &cli_key.pubkey, 10).await.unwrap(); let new_vm_resp = cli_client.new_vm(cli_key.sign_request(req).unwrap()).await.unwrap().into_inner();