Kick Contract #2

Merged
ghe0 merged 6 commits from general_and_admin into main 2025-05-23 13:14:54 +00:00
5 changed files with 41 additions and 20 deletions
Showing only changes of commit 540e5bfa4c - Show all commits

@ -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<Vec<Key>> = OnceLock::new();
pub fn admin_keys() -> Vec<Key> {
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,

@ -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?;

@ -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");

@ -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<ActiveVm> = 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();

@ -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();