use common::{ prepare_test_env::{ connect_stream_client_channel, prepare_test_db, run_service_for_stream_server, run_service_in_background, }, test_utils::Key, }; 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::RegisterVmNodeReq; use futures::StreamExt; use tokio::task::JoinSet; use tokio_stream::wrappers::ReceiverStream; mod common; #[tokio::test] async fn test_reg_vm_node() { prepare_test_db().await; let addr = run_service_in_background().await; let client = BrainVmDaemonClient::connect(format!("http://{}", addr)).await.unwrap(); let operator_wallet = Key::new().pubkey; let key = Key::new(); let vm_contracts = register_vm_node(client, key, operator_wallet).await; assert!(vm_contracts.is_empty()) } async fn register_vm_node( mut client: BrainVmDaemonClient, key: Key, operator_wallet: String, ) -> Vec { 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:?}"); } } } vm_contracts } #[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 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(); register_vm_node(daemon_client.clone(), daemon_key.clone(), Key::new().pubkey).await; 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)) => { todo!() } Some(vm_proto::brain_vm_message::Msg::DeleteVm(_del_vm_req)) => { todo!() } 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()); }