fix token calculation

all token inputs are in nano lp
fix MIN_ESCROW calculation
fixed all tests with TOKEN_DECIMAL multiplication

src/constants.rs

@@ -46,5 +46,5 @@ pub const ID_ALPHABET: [char; 62] = [
     'V', 'W', 'X', 'Y', 'Z',
 ];
 
-pub const MIN_ESCROW: u64 = 5000;
 pub const TOKEN_DECIMAL: u64 = 1_000_000_000;
+pub const MIN_ESCROW: u64 = 5000 * TOKEN_DECIMAL;

src/db/general.rs

@@ -1,5 +1,5 @@
 use super::Error;
-use crate::constants::{ACCOUNT, BAN, KICK, MIN_ESCROW, TOKEN_DECIMAL, VM_NODE};
+use crate::constants::{ACCOUNT, BAN, KICK, MIN_ESCROW, VM_NODE};
 use crate::db::prelude::*;
 use crate::old_brain;
 use serde::{Deserialize, Serialize};
@@ -42,7 +42,6 @@ impl Account {
     }
 
     pub async fn airdrop(db: &Surreal<Client>, account: &str, tokens: u64) -> Result<(), Error> {
-        let tokens = tokens.saturating_mul(1_000_000_000);
         let _ = db
             .query(format!("upsert account:{account} SET balance = (balance || 0) + {tokens};"))
             .await?;
@@ -64,7 +63,6 @@ impl Account {
             return Err(Error::MinimalEscrow);
         }
         let mut op_account = Self::get(db, wallet).await?;
-        let escrow = escrow.saturating_mul(TOKEN_DECIMAL);
         let op_total_balance = op_account.balance.saturating_add(op_account.escrow);
         if op_total_balance < escrow {
             return Err(Error::InsufficientFunds);
@@ -95,7 +93,7 @@ impl Account {
         let mut query_resp = db
             .query(tx_query)
             .bind(("account_input", RecordId::from((ACCOUNT, account))))
-            .bind(("slash_amount", slash_amount.saturating_mul(TOKEN_DECIMAL)))
+            .bind(("slash_amount", slash_amount))
             .await?;
 
         log::trace!("query_resp: {query_resp:?}");
@@ -259,7 +257,6 @@ pub struct Report {
 }
 
 impl Report {
-    // TODO: test this functionality and remove this comment
     pub async fn create(
         db: &Surreal<Client>,
         from_account: RecordId,

src/db/mod.rs

@@ -49,6 +49,10 @@ pub enum Error {
     AlreadyBanned(String),
     #[error("Failed to slash operator {0}")]
     FailedToSlashOperator(String),
+    #[error("Transation Too Big {0}")]
+    TooBigTransaction(String),
+    #[error("Unknown: {0}")]
+    Unknown(String),
 }
 
 pub mod prelude {

src/db/vm.rs

@@ -191,62 +191,120 @@ impl NewVmReq {
     }
 
     pub async fn submit_error(db: &Surreal<Client>, id: &str, error: String) -> Result<(), Error> {
-        #[derive(Serialize)]
-        struct NewVmError {
-            error: String,
+        let tx_query = String::from(
+            "
+            BEGIN TRANSACTION;
+
+                LET $new_vm_req = $new_vm_req_input;
+                LET $error = $error_input;
+                LET $record = (select * from $new_vm_req)[0];
+                LET $admin = $record.in;
+
+                if $record == None {{
+                    THROW 'vm req not exist ' + <string>$new_vm_req
+                }};
+
+                UPDATE $new_vm_req SET error = $error;
+
+                UPDATE $admin SET tmp_locked -= $record.locked_nano;
+                UPDATE $admin SET balance += $record.locked_nano;
+
+            COMMIT TRANSACTION;",
+        );
+
+        log::trace!("submit_new_vm_err query: {tx_query}");
+
+        let mut query_resp = db
+            .query(tx_query)
+            .bind(("new_vm_req_input", RecordId::from((NEW_VM_REQ, id))))
+            .bind(("error_input", error))
+            .await?;
+
+        let query_err = query_resp.take_errors();
+        if !query_err.is_empty() {
+            log::trace!("errors in submit_new_vm_err: {query_err:?}");
+            let tx_fail_err_str =
+                String::from("The query was not executed due to a failed transaction");
+
+            if query_err.contains_key(&4) && query_err[&4].to_string() != tx_fail_err_str {
+                log::error!("vm req not exist: {}", query_err[&4]);
+                return Err(Error::ContractNotFound);
+            } else {
+                log::error!("Unknown error in submit_new_vm_err: {query_err:?}");
+                return Err(Error::Unknown("submit_new_vm_req".to_string()));
             }
-        let _: Option<Self> = db.update((NEW_VM_REQ, id)).merge(NewVmError { error }).await?;
+        }
+
         Ok(())
     }
 
     pub async fn submit(self, db: &Surreal<Client>) -> Result<(), Error> {
         let locked_nano = self.locked_nano;
-        let account = self.admin.key().to_string();
-        let vm_id = self.id.key().to_string();
-        let vm_node = self.vm_node.key().to_string();
-        // TODO: check for possible injection and maybe use .bind()
-        let query = format!(
-            "
+        let tx_query = format!("
             BEGIN TRANSACTION;
-            UPDATE account:{account} SET balance -= {locked_nano};
-            IF account:{account}.balance < 0 {{
+
+                LET $account = $account_input;
+                LET $new_vm_req = $new_vm_req_input;
+                LET $vm_node = $vm_node_input;
+                LET $hostname = $hostname_input;
+                LET $dtrfs_url = $dtrfs_url_input;
+                LET $dtrfs_sha = $dtrfs_sha_input;
+                LET $kernel_url = $kernel_url_input;
+                LET $kernel_sha = $kernel_sha_input;
+
+                UPDATE $account SET balance -= {locked_nano};
+                IF $account.balance < 0 {{
                     THROW 'Insufficient funds.'
                 }};
-            UPDATE account:{account} SET tmp_locked += {locked_nano};
+                UPDATE $account SET tmp_locked += {locked_nano};
                 RELATE
-                account:{account}
-                ->new_vm_req:{vm_id}
-                ->vm_node:{vm_node}
+                    $account
+                    ->$new_vm_req
+                    ->$vm_node
                 CONTENT {{
-                created_at: time::now(), hostname: '{}', vcpus: {}, memory_mb: {}, disk_size_gb: {},
-                extra_ports: {:?}, public_ipv4: {:?}, public_ipv6: {:?},
-                dtrfs_url: '{}', dtrfs_sha: '{}', kernel_url: '{}', kernel_sha: '{}',
+                    created_at: time::now(), hostname: $hostname, vcpus: {}, memory_mb: {}, disk_size_gb: {},
+                    extra_ports: {:?}, public_ipv4: {}, public_ipv6: {},
+                    dtrfs_url: $dtrfs_url, dtrfs_sha: $dtrfs_sha, kernel_url: $kernel_url, kernel_sha: $kernel_sha,
                     price_per_unit: {}, locked_nano: {locked_nano}, error: ''
                 }};
-            COMMIT TRANSACTION;
-        ",
-            self.hostname,
+            
+            COMMIT TRANSACTION;",
             self.vcpus,
             self.memory_mb,
             self.disk_size_gb,
             self.extra_ports,
             self.public_ipv4,
             self.public_ipv6,
-            self.dtrfs_url,
-            self.dtrfs_sha,
-            self.kernel_url,
-            self.kernel_sha,
             self.price_per_unit
         );
-        //let _: Vec<Self> = db.insert(NEW_VM_REQ).relation(self).await?;
-        let mut query_resp = db.query(query).await?;
-        let resp_err = query_resp.take_errors();
 
-        if let Some(surrealdb::Error::Api(surrealdb::error::Api::Query(tx_query_error))) =
-            resp_err.get(&1)
-        {
-            log::error!("Transaction error: {tx_query_error}");
+        log::trace!("submit_new_vm_req query: {tx_query}");
+
+        let mut query_resp = db
+            .query(tx_query)
+            .bind(("account_input", self.admin))
+            .bind(("new_vm_req_input", self.id))
+            .bind(("vm_node_input", self.vm_node))
+            .bind(("hostname_input", self.hostname))
+            .bind(("dtrfs_url_input", self.dtrfs_url))
+            .bind(("dtrfs_sha_input", self.dtrfs_sha))
+            .bind(("kernel_url_input", self.kernel_url))
+            .bind(("kernel_sha_input", self.kernel_sha))
+            .await?;
+
+        let query_err = query_resp.take_errors();
+        if !query_err.is_empty() {
+            log::trace!("errors in submit_new_vm_req: {query_err:?}");
+            let tx_fail_err_str =
+                String::from("The query was not executed due to a failed transaction");
+
+            if query_err.contains_key(&9) && query_err[&9].to_string() != tx_fail_err_str {
+                log::error!("Transaction error: {}", query_err[&9]);
                 return Err(Error::InsufficientFunds);
+            } else {
+                log::error!("Unknown error in submit_new_vm_req: {query_err:?}");
+                return Err(Error::Unknown("submit_new_vm_req".to_string()));
+            }
         }
 
         Ok(())
@@ -271,14 +329,18 @@ impl WrappedMeasurement {
             UPDATE_VM_REQ => UPDATE_VM_REQ,
             _ => NEW_VM_REQ,
         };
-        let mut resp = db
-            .query(format!("live select error from {table} where id = {NEW_VM_REQ}:{vm_id};"))
+
+        let mut args_stream = db
             .query(format!(
                 "live select * from measurement_args where id = measurement_args:{vm_id};"
             ))
-            .await?;
-        let mut error_stream = resp.stream::<Notification<ErrorFromTable>>(0)?;
-        let mut args_stream = resp.stream::<Notification<vm_proto::MeasurementArgs>>(1)?;
+            .await?
+            .stream::<Notification<vm_proto::MeasurementArgs>>(0)?;
+
+        let mut error_stream = db
+            .query(format!("live select error from {table} where id = {NEW_VM_REQ}:{vm_id};"))
+            .await?
+            .stream::<Notification<ErrorFromTable>>(0)?;
 
         let args: Option<vm_proto::MeasurementArgs> =
             db.delete(("measurement_args", vm_id)).await?;
@@ -492,6 +554,78 @@ impl ActiveVm {
             Ok(false)
         }
     }
+
+    pub async fn extend_time(
+        db: &Surreal<Client>,
+        id: &str,
+        admin: &str,
+        nano_lp: u64,
+    ) -> Result<(), Error> {
+        if nano_lp > 100_000_000_000_000 {
+            return Err(Error::TooBigTransaction(nano_lp.to_string()));
+        }
+
+        let tx_query = format!(
+            "
+            BEGIN TRANSACTION;
+
+                LET $contract = $contract_input;
+                LET $admin = $admin_input;
+                LET $lock_amt = {nano_lp};
+
+                if !record::exists($contract) {{
+                    THROW 'contract not exist ' + <string>$contract
+                }};
+                if $contract.in != $admin {{
+                    THROW 'Unauthorized'
+                }};
+                if $admin.balance + $contract.locked_nano < $lock_amt {{
+                    THROW 'InsufficientFunds'
+                }};
+
+                UPDATE $admin SET balance = $admin.balance + $contract.locked_nano - $lock_amt;
+                UPDATE $contract SET locked_nano = $lock_amt;
+
+            COMMIT TRANSACTION;
+            "
+        );
+
+        log::trace!("extend_time query: {tx_query}");
+
+        let mut query_res = db
+            .query(tx_query)
+            .bind(("contract_input", RecordId::from((ACTIVE_VM, id))))
+            .bind(("admin_input", RecordId::from((ACCOUNT, admin))))
+            .await?;
+
+        log::trace!("tx_query response: {query_res:?}");
+
+        let query_err = query_res.take_errors();
+        if !query_err.is_empty() {
+            let tx_fail_err_str =
+                String::from("The query was not executed due to a failed transaction");
+
+            if query_err.contains_key(&3) && query_err[&3].to_string() != tx_fail_err_str {
+                log::error!("contract not exist: {}", query_err[&3]);
+                return Err(Error::ContractNotFound);
+            }
+
+            if query_err.contains_key(&4) && query_err[&4].to_string() != tx_fail_err_str {
+                log::error!("Unauthorized: {}", query_err[&4]);
+                return Err(Error::AccessDenied);
+            }
+
+            if query_err.contains_key(&5) && query_err[&5].to_string() != tx_fail_err_str {
+                log::error!("InsufficientFunds: {}", query_err[&5]);
+                return Err(Error::InsufficientFunds);
+            }
+
+            log::error!("Unknown error in extend_time: {query_err:?}");
+
+            return Err(Error::Unknown("extend_time".to_string()));
+        }
+        Ok(())
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]

src/grpc/vm.rs

@@ -313,15 +313,26 @@ impl BrainVmCli for VmCliServer {
     }
 
     async fn extend_vm(&self, req: Request<ExtendVmReq>) -> Result<Response<Empty>, Status> {
-        let _req = check_sig_from_req(req)?;
-        todo!();
-        // match self
-        //     .data
-        //     .extend_vm_contract_time(&req.uuid, &req.admin_pubkey, req.locked_nano)
-        // {
-        //     Ok(()) => Ok(Response::new(Empty {})),
-        //     Err(e) => Err(Status::unknown(format!("Could not extend contract: {e}"))),
-        // }
+        let req = check_sig_from_req(req)?;
+        match db::ActiveVm::extend_time(&self.db, &req.uuid, &req.admin_pubkey, req.locked_nano)
+            .await
+        {
+            Ok(()) => Ok(Response::new(Empty {})),
+            Err(e)
+                if matches!(
+                    e,
+                    db::Error::ContractNotFound
+                        | db::Error::AccessDenied
+                        | db::Error::InsufficientFunds
+                ) =>
+            {
+                Err(Status::failed_precondition(e.to_string()))
+            }
+            Err(e) => {
+                log::error!("Error extending VM contract {}: {e}", &req.uuid);
+                Err(Status::unknown(format!("Could not extend contract: {e}")))
+            }
+        }
     }
 
     async fn delete_vm(&self, req: Request<DeleteVmReq>) -> Result<Response<Empty>, Status> {

tests/common/vm_cli_utils.rs

@@ -7,7 +7,7 @@ use detee_shared::general_proto::{Account, AirdropReq, RegOperatorReq, ReportNod
 use detee_shared::vm_proto;
 use detee_shared::vm_proto::brain_vm_cli_client::BrainVmCliClient;
 use futures::StreamExt;
-use surreal_brain::constants::{ACTIVE_VM, NEW_VM_REQ};
+use surreal_brain::constants::{ACTIVE_VM, NEW_VM_REQ, TOKEN_DECIMAL};
 use surreal_brain::db::prelude as db;
 use surrealdb::engine::remote::ws::Client;
 use surrealdb::Surreal;
@@ -15,7 +15,7 @@ use tonic::transport::Channel;
 
 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 airdrop_req = AirdropReq { pubkey: wallet.to_string(), tokens: amount * TOKEN_DECIMAL };
 
     let admin_key = admin_keys()[0].clone();
 
@@ -46,7 +46,7 @@ pub async fn create_new_vm(
     assert!(new_vm_resp.uuid.len() == 40);
 
     // wait for update db
-    tokio::time::sleep(tokio::time::Duration::from_millis(700)).await;
+    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
 
     let vm_req_db: Option<db::NewVmReq> = db.select((NEW_VM_REQ, new_vm_resp.uuid.clone())).await?;
 
@@ -80,10 +80,13 @@ pub async fn report_node(
     Ok(client_gen_cli.report_node(key.sign_request(report_req)?).await?)
 }
 
-pub async fn register_operator(brain_channel: &Channel, key: &Key, escrow: u64) -> Result<()> {
+pub async fn register_operator(brain_channel: &Channel, key: &Key, escrow_nano: u64) -> Result<()> {
     let mut cli_client = BrainGeneralCliClient::new(brain_channel.clone());
-    let reg_req =
-        RegOperatorReq { pubkey: key.pubkey.clone(), escrow, email: "foo@bar.com".to_string() };
+    let reg_req = RegOperatorReq {
+        pubkey: key.pubkey.clone(),
+        escrow: escrow_nano,
+        email: "foo@bar.com".to_string(),
+    };
 
     cli_client.register_operator(key.sign_request(reg_req.clone()).unwrap()).await?;
     Ok(())
@@ -161,3 +164,38 @@ pub async fn list_all_app_contracts(
 
     Ok(app_contracts)
 }
+
+pub async fn user_list_vm_contracts(
+    brain_channel: &Channel,
+    key: &Key,
+    as_operator: bool,
+    uuid: &str,
+) -> Result<Vec<vm_proto::VmContract>> {
+    let mut cli_client = BrainVmCliClient::new(brain_channel.clone());
+    let mut stream = cli_client
+        .list_vm_contracts(
+            key.sign_request(vm_proto::ListVmContractsReq {
+                wallet: key.pubkey.clone(),
+                as_operator,
+                uuid: uuid.to_string(),
+            })
+            .unwrap(),
+        )
+        .await?
+        .into_inner();
+
+    let mut vm_contracts = Vec::new();
+
+    while let Some(stream_data) = stream.next().await {
+        match stream_data {
+            Ok(vm_contract) => {
+                vm_contracts.push(vm_contract);
+            }
+            Err(e) => {
+                panic!("Error while listing vm_contracts: {e:?}");
+            }
+        }
+    }
+
+    Ok(vm_contracts)
+}

tests/common/vm_daemon_utils.rs

@@ -8,7 +8,10 @@ use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::transport::Channel;
 
-pub async fn mock_vm_daemon(brain_channel: &Channel) -> Result<String> {
+pub async fn mock_vm_daemon(
+    brain_channel: &Channel,
+    daemon_error: Option<String>,
+) -> Result<String> {
     let mut daemon_client = BrainVmDaemonClient::new(brain_channel.clone());
     let daemon_key = Key::new();
 
@@ -26,7 +29,7 @@ pub async fn mock_vm_daemon(brain_channel: &Channel) -> Result<String> {
         rx,
     ));
 
-    tokio::spawn(daemon_engine(daemon_msg_tx.clone(), brain_msg_rx));
+    tokio::spawn(daemon_engine(daemon_msg_tx.clone(), brain_msg_rx, daemon_error));
 
     tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
 
@@ -102,6 +105,7 @@ pub async fn daemon_msg_sender(
 pub async fn daemon_engine(
     tx: mpsc::Sender<vm_proto::VmDaemonMessage>,
     mut rx: mpsc::Receiver<vm_proto::BrainVmMessage>,
+    new_vm_err: Option<String>,
 ) -> Result<()> {
     log::info!("daemon engine vm_daemon");
     while let Some(brain_msg) = rx.recv().await {
@@ -120,7 +124,7 @@ pub async fn daemon_engine(
                 let new_vm_resp = vm_proto::NewVmResp {
                     uuid: new_vm_req.uuid.clone(),
                     args,
-                    error: String::new(),
+                    error: new_vm_err.clone().unwrap_or_default(),
                 };
 
                 let res_data = vm_proto::VmDaemonMessage {

tests/grpc_general_test.rs

@@ -53,7 +53,8 @@ async fn test_general_airdrop() {
     let user_01_key = Key::new();
     let user_01_pubkey = user_01_key.pubkey.clone();
 
-    let airdrop_req = AirdropReq { pubkey: user_01_pubkey.clone(), tokens: airdrop_amount };
+    let airdrop_req =
+        AirdropReq { pubkey: user_01_pubkey.clone(), tokens: airdrop_amount * TOKEN_DECIMAL };
 
     // user airdroping himself
     let err =
@@ -97,7 +98,8 @@ async fn test_general_airdrop() {
     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 };
+    let airdrop_req =
+        AirdropReq { pubkey: admin_keys[2].pubkey.clone(), tokens: airdrop_amount * TOKEN_DECIMAL };
 
     let _ = client.airdrop(admin_keys[2].sign_request(airdrop_req.clone()).unwrap()).await.unwrap();
 
@@ -116,7 +118,7 @@ async fn test_report_node() {
     let db = prepare_test_db().await.unwrap();
 
     let brain_channel = run_service_for_stream().await.unwrap();
-    let daemon_key = mock_vm_daemon(&brain_channel).await.unwrap();
+    let daemon_key = mock_vm_daemon(&brain_channel, None).await.unwrap();
 
     let key = Key::new();
 
@@ -223,17 +225,19 @@ async fn test_register_operator() {
     let min_escrew_error = register_operator(&brain_channel, &key, 10).await.err().unwrap();
     assert!(min_escrew_error.to_string().contains("Minimum escrow amount is 5000"));
 
-    let no_balance = register_operator(&brain_channel, &key, 5000).await.err().unwrap();
+    let no_balance =
+        register_operator(&brain_channel, &key, 5000 * TOKEN_DECIMAL).await.err().unwrap();
     assert!(no_balance.to_string().contains("Insufficient funds, deposit more tokens"));
 
     airdrop(&brain_channel, &key.pubkey, 1000).await.unwrap();
 
-    let no_balance = register_operator(&brain_channel, &key, 5000).await.err().unwrap();
+    let no_balance =
+        register_operator(&brain_channel, &key, 5000 * TOKEN_DECIMAL).await.err().unwrap();
     assert!(no_balance.to_string().contains("Insufficient funds, deposit more tokens"));
 
     airdrop(&brain_channel, &key.pubkey, 7000).await.unwrap();
 
-    register_operator(&brain_channel, &key, 6000).await.unwrap();
+    register_operator(&brain_channel, &key, 6000 * TOKEN_DECIMAL).await.unwrap();
 
     let operator_account: Option<db::Account> =
         db_conn.select((ACCOUNT, key.pubkey)).await.unwrap();
@@ -356,9 +360,9 @@ async fn test_slash_operator() {
     let slash_amt = 2500;
 
     airdrop(&brain_channel, &op_key.pubkey, 10000).await.unwrap();
-    register_operator(&brain_channel, &op_key, escrew).await.unwrap();
+    register_operator(&brain_channel, &op_key, escrew * TOKEN_DECIMAL).await.unwrap();
 
-    let raw_slash_req = SlashReq { pubkey: op_key.pubkey.clone(), tokens: 2500 };
+    let raw_slash_req = SlashReq { pubkey: op_key.pubkey.clone(), tokens: 2500 * TOKEN_DECIMAL };
 
     let other_key = Key::new();
     let admin_error = cli_client
@@ -399,9 +403,9 @@ async fn test_admin_list_account() {
 
     airdrop(&brain_channel, &Key::new().pubkey, 10).await.unwrap();
 
+    let accounts = list_accounts(&brain_channel, &admin_key).await.unwrap();
     let acc_in_db = db_conn.select::<Vec<db::Account>>(ACCOUNT).await.unwrap();
 
-    let accounts = list_accounts(&brain_channel, &admin_key).await.unwrap();
     assert_eq!(accounts.len(), acc_in_db.len());
 }
 

tests/grpc_vm_cli_test.rs

@@ -1,24 +1,31 @@
 use common::prepare_test_env::{prepare_test_db, run_service_for_stream};
 use common::test_utils::Key;
-use common::vm_cli_utils::{airdrop, create_new_vm};
+use common::vm_cli_utils::{airdrop, 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::{ListVmContractsReq, NewVmReq};
+use detee_shared::vm_proto::{ExtendVmReq, ListVmContractsReq, NewVmReq};
 use futures::StreamExt;
 use std::vec;
-use surreal_brain::constants::ACTIVE_VM;
-use surreal_brain::db::vm::ActiveVm;
+use surreal_brain::constants::{ACCOUNT, ACTIVE_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::Error).init();
+    /*
+    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).await.unwrap();
+    let daemon_key = mock_vm_daemon(&brain_channel, None).await.unwrap();
 
     let key = Key::new();
 
@@ -28,11 +35,60 @@ async fn test_vm_creation() {
     let grpc_error_message = new_vm_resp.err().unwrap().to_string();
     assert!(grpc_error_message.contains("Insufficient funds"));
 
-    airdrop(&brain_channel, &key.pubkey, 10).await.unwrap();
+    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<ActiveVm> = db.select((ACTIVE_VM, new_vm_id)).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);
+    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);
+    assert_eq!(acc_db.tmp_locked, 0);
 }
 
 #[tokio::test]
@@ -102,3 +158,85 @@ async fn test_list_vm_contracts() {
 
     // 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);
+}

tests/grpc_vm_daemon_test.rs

@@ -29,7 +29,7 @@ async fn test_brain_message() {
     prepare_test_db().await.unwrap();
 
     let brain_channel = run_service_for_stream().await.unwrap();
-    let daemon_key = mock_vm_daemon(&brain_channel).await.unwrap();
+    let daemon_key = mock_vm_daemon(&brain_channel, None).await.unwrap();
     let mut cli_client = BrainVmCliClient::new(brain_channel.clone());
 
     let cli_key = Key::new();

tests/mock_data.yaml

@@ -114,6 +114,18 @@ accounts:
     kicked_for: []
     last_kick: 1970-01-01T00:00:00Z
     banned_by: []
+  Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG:
+    balance: 25949200000
+    tmp_locked: 0
+    kicked_for: []
+    last_kick: 1970-01-01T00:00:00Z
+    banned_by: []
+  GmE4JH3bL4NpmzwKCBJemJzRTumJAnbcXLGqce5mREgS:
+    balance: 500000000000
+    tmp_locked: 0
+    kicked_for: []
+    last_kick: 1970-01-01T00:00:00Z
+    banned_by: []
 
 operators:
   BFopWmwcZAMF1h2PFECZNdEucdZfnZZ32p6R9ZaBiVsS:
@@ -151,6 +163,14 @@ operators:
       - 7fujZQeTme52RdXTLmQST5jBgAbvzic5iERtH5EWoYjk
     app_nodes: []
 
+  Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG:
+    escrow: 5499700480000
+    email: "test_mock_extend@operator"
+    banned_users: []
+    vm_nodes:
+      - 8ue3VHMnJg2i8pwTQ6mJtvYuS2kd9n1XLLco8GUPfT95
+    app_nodes: []
+
 vm_nodes:
   - public_key: 7Xw3RxbP5pvfjZ8U6yA3HHVSS9YXjKH5Vkas3JRbQYd9
     operator_wallet: x52w7jARC5erhWWK65VZmjdGXzBK6ZDgfv1A283d8XK
@@ -265,6 +285,22 @@ vm_nodes:
     price: 24000
     reports: {}
     offline_minutes: 0
+  - public_key: 8ue3VHMnJg2i8pwTQ6mJtvYuS2kd9n1XLLco8GUPfT95
+    operator_wallet: Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG
+    country: GB
+    region: England
+    city: London
+    ip: 193.234.17.2
+    avail_mem_mb: 28000
+    avail_vcpus: 24
+    avail_storage_gbs: 1680
+    avail_ipv4: 1
+    avail_ipv6: 0
+    avail_ports: 19999
+    max_ports_per_vm: 10
+    price: 24000
+    reports: {}
+    offline_minutes: 0
 
 vm_contracts:
   - uuid: 958165e3-dea8-407d-8c42-dd17002ef79c
@@ -406,6 +442,24 @@ vm_contracts:
     locked_nano: 12730960000
     collected_at: 2025-04-20T00:34:15.461240342Z
 
+  - uuid: 1d749a81-6c27-a22e-fa0e574-b023-a6afef040fe5
+    hostname: test-extend
+    admin_pubkey: GmE4JH3bL4NpmzwKCBJemJzRTumJAnbcXLGqce5mREgS
+    node_pubkey: 8ue3VHMnJg2i8pwTQ6mJtvYuS2kd9n1XLLco8GUPfT95
+    exposed_ports:
+      - 46393
+    public_ipv4: ""
+    public_ipv6: ""
+    disk_size_gb: 10
+    vcpus: 1
+    memory_mb: 1000
+    kernel_sha: e765e56166ef321b53399b9638584d1279821dbe3d46191c1f66bbaa075e7919
+    dtrfs_sha: d207644ee60d54009b6ecdfb720e2ec251cde31774dd249fcc7435aca0377990
+    created_at: 2025-04-16T20:37:57.176592933Z
+    updated_at: 2025-04-16T20:37:57.176594069Z
+    price_per_unit: 20000
+    locked_nano: 60000
+    collected_at: 2025-04-20T00:34:15.461240342Z
   - uuid: 5af49a71-4c64-a82e-f50e574-b023-b2a0ef0405ed
     hostname: hallow-hobo
     admin_pubkey: 4qFJJJdRrSB9hCn8rrvYTXHLJg371ab36PJmZ4uxHjGQ