Add tests for Burn API

contract/src/auth/tests.rs

@@ -1,127 +1,70 @@
 #![cfg(test)]
 
-use model::api::{AuthApi, InitApi};
-use near_sdk::{test_utils::VMContextBuilder, testing_env, AccountId};
+use model::api::AuthApi;
 
-use crate::Contract;
+use crate::common::tests::Context;
 
 #[test]
 fn add_oracle_by_contract_owner() {
-    Context::run_test(|mut context, mut contract, accounts| {
-        context.switch_account(&accounts.owner);
-        contract.add_oracle(accounts.oracle.clone());
+    let (mut context, mut contract, accounts) = Context::init();
+    context.switch_account(&accounts.owner);
+    contract.add_oracle(accounts.oracle.clone());
 
-        let oracles = contract.get_oracles();
-        assert_eq!(oracles, vec![accounts.oracle.clone()]);
-    });
+    let oracles = contract.get_oracles();
+    assert_eq!(oracles, vec![accounts.oracle.clone()]);
 }
 
 #[test]
 #[should_panic(expected = "Method is private")]
 fn add_oracle_not_by_contract_owner() {
-    Context::run_test(|mut context, mut contract, accounts| {
-        context.switch_account(&accounts.alice);
-        contract.add_oracle(accounts.oracle.clone());
-    });
+    let (mut context, mut contract, accounts) = Context::init();
+
+    context.switch_account(&accounts.alice);
+    contract.add_oracle(accounts.oracle.clone());
 }
 
 #[test]
 #[should_panic(expected = "Already exists")]
 fn add_oracle_twice() {
-    Context::run_test(|mut context, mut contract, accounts| {
-        context.switch_account(&accounts.owner);
-        contract.add_oracle(accounts.oracle.clone());
-        contract.add_oracle(accounts.oracle.clone());
-    });
+    let (mut context, mut contract, accounts) = Context::init();
+
+    context.switch_account(&accounts.owner);
+    contract.add_oracle(accounts.oracle.clone());
+    contract.add_oracle(accounts.oracle.clone());
 }
 
 #[test]
 fn remove_oracle_by_contract_owner() {
-    Context::run_test(|mut context, mut contract, accounts| {
-        context.switch_account(&accounts.owner);
-        contract.add_oracle(accounts.oracle.clone());
+    let (mut context, mut contract, accounts) = Context::init();
+
+    context.switch_account(&accounts.owner);
+    contract.add_oracle(accounts.oracle.clone());
 
-        let oracles = contract.get_oracles();
-        assert_eq!(oracles, vec![accounts.oracle.clone()]);
+    let oracles = contract.get_oracles();
+    assert_eq!(oracles, vec![accounts.oracle.clone()]);
 
-        contract.remove_oracle(accounts.oracle.clone());
+    contract.remove_oracle(accounts.oracle.clone());
 
-        let oracles = contract.get_oracles();
-        assert!(oracles.is_empty());
-    });
+    let oracles = contract.get_oracles();
+    assert!(oracles.is_empty());
 }
 
 #[test]
 #[should_panic(expected = "Method is private")]
 fn remove_oracle_not_by_contract_owner() {
-    Context::run_test(|mut context, mut contract, accounts| {
-        contract.oracles.insert(accounts.oracle.clone());
+    let (mut context, mut contract, accounts) = Context::init();
 
-        context.switch_account(&accounts.alice);
-        contract.remove_oracle(accounts.oracle.clone());
-    });
+    contract.oracles.insert(accounts.oracle.clone());
+
+    context.switch_account(&accounts.alice);
+    contract.remove_oracle(accounts.oracle.clone());
 }
 
 #[test]
 #[should_panic(expected = "No such oracle")]
 fn remove_not_existing_oracle() {
-    Context::run_test(|mut context, mut contract, accounts| {
-        context.switch_account(&accounts.owner);
-        contract.remove_oracle(accounts.oracle.clone());
-    });
-}
-
-struct Context {
-    builder: VMContextBuilder,
-}
-
-impl Context {
-    fn run_test<F>(test: F)
-    where
-        F: Fn(Context, Contract, &TestAccounts),
-    {
-        let accounts = TestAccounts::default();
-        let token_account = accounts.token.clone();
-
-        let mut builder = VMContextBuilder::new();
-        builder
-            .current_account_id(accounts.owner.clone())
-            .signer_account_id(accounts.owner.clone())
-            .predecessor_account_id(accounts.owner.clone())
-            .block_timestamp(0);
-
-        testing_env!(builder.build());
-
-        let contract = Contract::init(token_account);
-        let context = Context { builder };
-
-        test(context, contract, &accounts);
-    }
-
-    pub(crate) fn switch_account(&mut self, account_id: &AccountId) {
-        self.builder
-            .predecessor_account_id(account_id.clone())
-            .signer_account_id(account_id.clone());
-        testing_env!(self.builder.build());
-    }
-}
-
-struct TestAccounts {
-    pub alice: AccountId,
-    pub bob: AccountId,
-    pub oracle: AccountId,
-    pub token: AccountId,
-    pub owner: AccountId,
-}
+    let (mut context, mut contract, accounts) = Context::init();
 
-impl Default for TestAccounts {
-    fn default() -> Self {
-        Self {
-            alice: AccountId::new_unchecked("alice".to_string()),
-            bob: AccountId::new_unchecked("bob".to_string()),
-            oracle: AccountId::new_unchecked("oracle".to_string()),
-            token: AccountId::new_unchecked("token".to_string()),
-            owner: AccountId::new_unchecked("owner".to_string()),
-        }
-    }
+    context.switch_account(&accounts.owner);
+    contract.remove_oracle(accounts.oracle.clone());
 }

contract/src/burn/api.rs

@@ -1,16 +1,25 @@
 use model::{api::BurnApi, TokensAmount, UnixTimestamp};
 use near_sdk::{
-    env, ext_contract, is_promise_success, json_types::U128, near_bindgen, serde_json::json, Gas, Promise,
+    env, ext_contract, is_promise_success, json_types::U128, near_bindgen, require, serde_json::json, Gas, Promise,
     PromiseOrValue,
 };
 
+#[cfg(test)]
+use crate::common::tests::data::get_test_future_success;
 use crate::{common::unix_timestamp, Contract, ContractExt};
 
+#[cfg(test)]
+pub(crate) const EXT_BURN_FUTURE: &str = "ext_burn";
+
 #[near_bindgen]
 impl BurnApi for Contract {
     fn burn(&mut self) -> PromiseOrValue<U128> {
         self.assert_oracle();
 
+        require!(!self.is_service_call_running, "Another service call is running");
+
+        self.is_service_call_running = true;
+
         let mut total_to_burn = 0;
         let mut keys_to_remove: Vec<UnixTimestamp> = vec![];
         let now: UnixTimestamp = unix_timestamp(env::block_timestamp_ms());
@@ -22,6 +31,37 @@ impl BurnApi for Contract {
             }
         }
 
+        if total_to_burn > 0 {
+            self.burn_external(total_to_burn, keys_to_remove)
+        } else {
+            self.is_service_call_running = false;
+
+            PromiseOrValue::Value(U128(0))
+        }
+    }
+}
+
+#[cfg(not(test))]
+#[ext_contract(ext_self)]
+pub trait SelfCallback {
+    fn on_burn(&mut self, total_to_burn: TokensAmount, keys_to_remove: Vec<UnixTimestamp>) -> U128;
+}
+
+#[cfg(not(test))]
+#[near_bindgen]
+impl SelfCallback for Contract {
+    fn on_burn(&mut self, total_to_burn: TokensAmount, keys_to_remove: Vec<UnixTimestamp>) -> U128 {
+        self.on_burn_internal(total_to_burn, keys_to_remove, is_promise_success())
+    }
+}
+
+impl Contract {
+    #[cfg(not(test))]
+    fn burn_external(
+        &mut self,
+        total_to_burn: TokensAmount,
+        keys_to_remove: Vec<UnixTimestamp>,
+    ) -> PromiseOrValue<U128> {
         let args = json!({
             "amount": U128(total_to_burn),
         })
@@ -38,17 +78,29 @@ impl BurnApi for Contract {
             )
             .into()
     }
-}
 
-#[ext_contract(ext_self)]
-pub trait SelfCallback {
-    fn on_burn(&mut self, total_to_burn: TokensAmount, keys_to_remove: Vec<UnixTimestamp>) -> U128;
-}
+    #[cfg(test)]
+    fn burn_external(
+        &mut self,
+        total_to_burn: TokensAmount,
+        keys_to_remove: Vec<UnixTimestamp>,
+    ) -> PromiseOrValue<U128> {
+        PromiseOrValue::Value(self.on_burn_internal(
+            total_to_burn,
+            keys_to_remove,
+            get_test_future_success(EXT_BURN_FUTURE),
+        ))
+    }
 
-#[near_bindgen]
-impl SelfCallback for Contract {
-    fn on_burn(&mut self, total_to_burn: TokensAmount, keys_to_remove: Vec<UnixTimestamp>) -> U128 {
-        if is_promise_success() {
+    fn on_burn_internal(
+        &mut self,
+        total_to_burn: TokensAmount,
+        keys_to_remove: Vec<UnixTimestamp>,
+        is_success: bool,
+    ) -> U128 {
+        self.is_service_call_running = false;
+
+        if is_success {
             for datetime in keys_to_remove {
                 self.accruals.remove(&datetime);
             }

contract/src/burn/mod.rs

@@ -1 +1,2 @@
 pub(crate) mod api;
+mod tests;

contract/src/burn/tests.rs

@@ -0,0 +1,106 @@
+#![cfg(test)]
+
+use model::api::{BurnApi, ClaimApi, RecordApi};
+use near_sdk::{json_types::U128, PromiseOrValue};
+
+use crate::{
+    burn::api::EXT_BURN_FUTURE,
+    common::tests::{data::set_test_future_success, Context},
+};
+
+#[test]
+fn test_burn_when_outdated_tokens_exist() {
+    let (mut context, mut contract, accounts) = Context::init_with_oracle();
+    set_test_future_success(EXT_BURN_FUTURE, true);
+
+    let alice_balance = 100_000;
+    let bob_balance = 200_000;
+
+    context.switch_account(&accounts.oracle);
+    contract.record_batch_for_hold(vec![
+        (accounts.alice.clone(), U128(alice_balance)),
+        (accounts.bob.clone(), U128(bob_balance)),
+    ]);
+
+    context.set_block_timestamp_in_seconds(contract.burn_period as u64 + 100);
+
+    let burn_result = contract.burn();
+    let burnt_amount = match burn_result {
+        PromiseOrValue::Promise(_) => panic!("Expected value"),
+        PromiseOrValue::Value(value) => value.0,
+    };
+
+    assert_eq!(alice_balance + bob_balance, burnt_amount);
+
+    let alice_new_balance = contract.get_claimable_balance_for_account(accounts.alice).0;
+    assert_eq!(0, alice_new_balance);
+
+    let bob_new_balance = contract.get_claimable_balance_for_account(accounts.bob).0;
+    assert_eq!(0, bob_new_balance);
+
+    assert!(!contract.is_service_call_running);
+}
+
+#[test]
+fn test_ext_error_on_burn_when_outdated_tokens_exist() {
+    let (mut context, mut contract, accounts) = Context::init_with_oracle();
+    set_test_future_success(EXT_BURN_FUTURE, false);
+
+    let alice_balance = 100_000;
+    let bob_balance = 200_000;
+
+    context.switch_account(&accounts.oracle);
+    contract.record_batch_for_hold(vec![
+        (accounts.alice.clone(), U128(alice_balance)),
+        (accounts.bob.clone(), U128(bob_balance)),
+    ]);
+
+    context.set_block_timestamp_in_seconds(contract.burn_period as u64 + 100);
+
+    let burn_result = contract.burn();
+    let burnt_amount = match burn_result {
+        PromiseOrValue::Promise(_) => panic!("Expected value"),
+        PromiseOrValue::Value(value) => value.0,
+    };
+
+    assert_eq!(0, burnt_amount);
+
+    let alice_new_balance = contract.get_claimable_balance_for_account(accounts.alice).0;
+    assert_eq!(0, alice_new_balance);
+
+    let bob_new_balance = contract.get_claimable_balance_for_account(accounts.bob).0;
+    assert_eq!(0, bob_new_balance);
+
+    assert!(!contract.is_service_call_running);
+}
+
+#[test]
+fn test_burn_when_outdated_tokens_don_not_exist() {
+    let (mut context, mut contract, accounts) = Context::init_with_oracle();
+    set_test_future_success(EXT_BURN_FUTURE, true);
+
+    let alice_balance = 500_000;
+    let bob_balance = 300_000;
+
+    context.switch_account(&accounts.oracle);
+    contract.record_batch_for_hold(vec![
+        (accounts.alice.clone(), U128(alice_balance)),
+        (accounts.bob.clone(), U128(bob_balance)),
+    ]);
+
+    let burn_result = contract.burn();
+    let burnt_amount = match burn_result {
+        PromiseOrValue::Promise(_) => panic!("Expected value"),
+        PromiseOrValue::Value(value) => value.0,
+    };
+
+    assert_eq!(0, burnt_amount);
+
+    let alice_new_balance = contract.get_claimable_balance_for_account(accounts.alice).0;
+    assert_eq!(alice_balance, alice_new_balance);
+
+    let bob_new_balance = contract.get_claimable_balance_for_account(accounts.bob).0;
+    assert_eq!(bob_balance, bob_new_balance);
+
+    assert!(!contract.is_service_call_running);
+}

contract/src/common/mod.rs

@@ -2,6 +2,7 @@ use model::UnixTimestamp;
 use near_sdk::env::panic_str;
 
 mod asserts;
+pub(crate) mod tests;
 
 pub(crate) fn unix_timestamp(ms: u64) -> UnixTimestamp {
     u32::try_from(ms / 1000)