use detee_shared::common_proto::Empty; use detee_shared::general_proto::ReportNodeReq; 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, RegisterVmNodeReq}; use detee_shared::{ common_proto::Pubkey, general_proto::brain_general_cli_client::BrainGeneralCliClient, vm_proto, }; mod common; use common::prepare_test_env::{ connect_stream_client_channel, prepare_test_db, run_service_for_stream, run_service_for_stream_server, run_service_in_background, }; use common::test_utils::Key; use futures::StreamExt; use tokio::task::JoinSet; use tokio_stream::wrappers::ReceiverStream; #[tokio::test] async fn test_reg_vm_node() { prepare_test_db().await; let addr = run_service_in_background().await; let mut client = BrainVmDaemonClient::connect(format!("http://{}", addr)).await.unwrap(); let operator_wallet = Key::new().pubkey; let key = Key::new(); let node_pubkey = key.pubkey.clone(); let req = RegisterVmNodeReq { node_pubkey, operator_wallet, main_ip: String::from("185.243.218.213"), city: String::from("Oslo"), country: String::from("Norway"), region: String::from("EU"), price: 1200, }; let mut grpc_stream = client.register_vm_node(key.sign_request(req).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()) } #[tokio::test] async fn test_brain_message() { // spawn grpc stream server // mock a daemon // mock a cli client to interact with brain // validate if something happening in "surreal_brain::db::NewVmReq", "surreal_brain::db::UpdateVmReq", "surreal_brain::db::DeletedVm" these table // mock daemon will responde to brain // TODO: register node before connecting to brain env_logger::builder().filter_level(log::LevelFilter::Info).init(); let _ = prepare_test_db().await; let (tokio_duplex, addr) = run_service_for_stream_server().await; let channel = connect_stream_client_channel(tokio_duplex, addr).await; let mut daemon_client = BrainVmDaemonClient::new(channel.clone()); let daemon_key = Key::new(); let mut daemon_join_set = JoinSet::new(); let (tx, mut brain_msg_rx) = tokio::sync::mpsc::channel(1); // listen to brain let mut daemon_client_01 = daemon_client.clone(); let daemon_key_01 = daemon_key.clone(); daemon_join_set.spawn(async move { let mut grpc_stream = daemon_client_01 .brain_messages(daemon_key_01.sign_stream_auth(vec![]).unwrap()) .await .unwrap() .into_inner(); while let Some(Ok(stream_update)) = grpc_stream.next().await { log::info!("vm deamon got notified: {:?}", &stream_update); let _ = tx.send(stream_update).await; } }); // send to brain let (daemon_msg_tx, rx) = tokio::sync::mpsc::channel(1); let daemon_msg_tx_01 = daemon_msg_tx.clone(); let daemon_key_02 = daemon_key.clone(); daemon_join_set.spawn(async move { let rx_stream = ReceiverStream::new(rx); daemon_msg_tx_01 .send(vm_proto::VmDaemonMessage { msg: Some(vm_proto::vm_daemon_message::Msg::Auth( daemon_key_02.sign_stream_auth(vec![]).unwrap(), )), }) .await .unwrap(); daemon_client.daemon_messages(rx_stream).await.unwrap(); }); // daemon engine daemon_join_set.spawn(async move { while let Some(brain_msg) = brain_msg_rx.recv().await { match brain_msg.msg { Some(vm_proto::brain_vm_message::Msg::NewVmReq(new_vm_req)) => { let args = Some(vm_proto::MeasurementArgs { dtrfs_api_endpoint: String::from("184.107.169.199:48865"), exposed_ports: new_vm_req.extra_ports, ovmf_hash: String::from( "0346619257269b9a61ee003e197d521b8e2283483070d163a34940d6a1d40d76", ), ips: vec![], }); let new_vm_resp = vm_proto::NewVmResp { uuid: new_vm_req.uuid.clone(), args, error: String::new(), }; let res_data = vm_proto::VmDaemonMessage { msg: Some(vm_proto::vm_daemon_message::Msg::NewVmResp(new_vm_resp)), }; daemon_msg_tx.send(res_data).await.unwrap(); } Some(vm_proto::brain_vm_message::Msg::UpdateVmReq(_update_vm_req)) => { // } Some(vm_proto::brain_vm_message::Msg::DeleteVm(_del_vm_req)) => { // } None => todo!(), } } }); tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; let mut cli_client = BrainVmCliClient::new(channel); let cli_key = Key::new(); let req = vm_proto::NewVmReq { admin_pubkey: cli_key.pubkey.clone(), node_pubkey: daemon_key.pubkey.clone(), price_per_unit: 1200, extra_ports: vec![8080, 8081], locked_nano: 0, ..Default::default() }; let new_vm_resp = cli_client.new_vm(cli_key.sign_request(req).unwrap()).await.unwrap().into_inner(); assert!(new_vm_resp.error.is_empty()); assert!(new_vm_resp.uuid.len() == 40); let id = ("measurement_args", new_vm_resp.uuid); let data_in_db: detee_shared::vm_proto::MeasurementArgs = surreal_brain::db::DB.select(id).await.unwrap().unwrap(); assert_eq!(data_in_db, new_vm_resp.args.unwrap()); } #[tokio::test] async fn test_general_balance() { // env_logger::builder().filter_level(log::LevelFilter::Trace).init(); let _ = prepare_test_db().await; let addr = run_service_in_background().await; let mut client = BrainGeneralCliClient::connect(format!("http://{}", addr)).await.unwrap(); let key = Key::new(); let pubkey = key.pubkey.clone(); let req_data = Pubkey { pubkey }; let req = key.sign_request(req_data).unwrap(); let acc_bal = client.get_balance(req).await.unwrap().into_inner(); assert_eq!(acc_bal.balance, 0); assert_eq!(acc_bal.tmp_locked, 0); } #[tokio::test] async fn test_report_node() { prepare_test_db().await; let addr = run_service_in_background().await; let mut client = BrainGeneralCliClient::connect(format!("http://{}", addr)).await.unwrap(); let key = Key::new(); let pubkey = key.pubkey.clone(); // TODO: create contract, node and operator in db and use it here let req_data = ReportNodeReq { admin_pubkey: pubkey, node_pubkey: String::from("node_pubkey"), contract: String::from("uuid"), reason: String::from("reason"), }; let report_error = client.report_node(key.sign_request(req_data).unwrap()).await.err().unwrap(); println!("Report error: {:?}", report_error); assert_eq!(report_error.message(), "No contract found by this ID."); // verify report in db } #[tokio::test] // TODO: register some operators before testing this async fn test_list_operators() { prepare_test_db().await; let channel = run_service_for_stream().await; let mut client = BrainGeneralCliClient::new(channel); let key = Key::new(); let mut grpc_stream = client.list_operators(key.sign_request(Empty {}).unwrap()).await.unwrap().into_inner(); let mut operators = Vec::new(); while let Some(stream_update) = grpc_stream.next().await { match stream_update { Ok(op) => { operators.push(op); } Err(e) => { panic!("Received error instead of operators: {e:?}"); } } } assert!(!operators.is_empty()) } #[tokio::test] // TODO: create vm for this user before testing this async fn test_list_vm_contracts() { prepare_test_db().await; let channel = run_service_for_stream().await; 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 }