fix: improved withdraw all and tests (#84)

* fix: improved withdraw all and tests

* fix: integration tests

* fix: added typos check, fixed typos, fixed mutation tests

* fix: review fixes

* chore: measured restake_all gas

* feat: extensive gas_left tests

.github/workflows/push.yml

@@ -8,6 +8,11 @@ env:
   CARGO_TERM_COLOR: always
 
 jobs:
+  typos:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: crate-ci/[email protected]
 
   lint:
     runs-on: ubuntu-latest
@@ -80,7 +85,7 @@ jobs:
           path: measured.txt
 
   push:
-    needs: [ build, lint, unit-tests, integration-tests, measure, mutation-tests ]
+    needs: [ build, typos, lint, unit-tests, integration-tests, measure, mutation-tests ]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout

.github/workflows/test.yml

@@ -8,6 +8,11 @@ env:
   CARGO_TERM_COLOR: always
 
 jobs:
+  typos:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: crate-ci/[email protected]
 
   lint:
     runs-on: ubuntu-latest

Cargo.toml

@@ -7,7 +7,6 @@ members = ["model", "contract", "integration-tests"]
 
 [workspace.dependencies]
 anyhow = "1.0.75"
-async-trait = "0.1.74"
 fake = "2.8.0"
 rand = "0.8.5"
 futures = "0.3.28"
@@ -16,8 +15,8 @@ ed25519-dalek = { version = "2.0.0", features = ["rand_core"] }
 base64 = "0.22.0"
 sha256 = "1.3.0"
 mutants = "0.0.3"
-serde = "1.0"
 sha2 = "0.10"
+serde = "1.0"
 
 nitka = "0.4.0"
 nitka-proc = "0.4.0"

contract/src/integration_test/integration_test.rs

@@ -1,14 +1,50 @@
 #![cfg(feature = "integration-test")]
 
-use near_sdk::{env, near_bindgen, Timestamp};
-use sweat_jar_model::api::IntegrationTestMethods;
+use near_sdk::{env, near_bindgen, AccountId, Timestamp};
+use sweat_jar_model::{api::IntegrationTestMethods, jar::JarView, ProductId};
 
-use crate::{Contract, ContractExt};
+use crate::{jar::model::Jar, Contract, ContractExt};
 
+#[mutants::skip]
 #[near_bindgen]
 impl IntegrationTestMethods for Contract {
-    #[mutants::skip]
     fn block_timestamp_ms(&self) -> Timestamp {
         env::block_timestamp_ms()
     }
+
+    fn bulk_create_jars(
+        &mut self,
+        account_id: AccountId,
+        product_id: ProductId,
+        principal: u128,
+        number_of_jars: u16,
+    ) -> Vec<JarView> {
+        self.assert_manager();
+        (0..number_of_jars)
+            .map(|_| self.create_jar_for_integration_tests(&account_id, &product_id, principal))
+            .collect()
+    }
+}
+
+#[mutants::skip]
+impl Contract {
+    fn create_jar_for_integration_tests(
+        &mut self,
+        account_id: &AccountId,
+        product_id: &ProductId,
+        amount: u128,
+    ) -> JarView {
+        let product = self.get_product(&product_id);
+
+        product.assert_enabled();
+        product.assert_cap(amount);
+
+        let id = self.increment_and_get_last_jar_id();
+        let now = env::block_timestamp_ms();
+        let jar = Jar::create(id, account_id.clone(), product_id.clone(), amount, now);
+
+        self.add_new_jar(account_id, jar.clone());
+
+        jar.into()
+    }
 }

contract/src/internal.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, fmt::Display};
 
 use near_sdk::require;
 use sweat_jar_model::{
@@ -74,10 +74,30 @@ impl Contract {
     }
 }
 
+pub(crate) fn assert_gas<Message: Display>(gas_needed: u64, error: impl FnOnce() -> Message) {
+    let gas_left = env::prepaid_gas().as_gas() - env::used_gas().as_gas();
+
+    if gas_left < gas_needed {
+        let error = error();
+
+        env::panic_str(&format!(
+            r#"Not enough gas left. Consider attaching more gas to the transaction.
+               {error}
+               Gas left: {gas_left} Needed: {gas_needed}. Need additional {} gas"#,
+            gas_needed - gas_left
+        ));
+    }
+}
+
 #[cfg(test)]
 mod test {
+    use near_sdk::env;
 
-    use crate::{common::tests::Context, test_utils::admin};
+    use crate::{
+        common::tests::Context,
+        internal::assert_gas,
+        test_utils::{admin, expect_panic},
+    };
 
     #[test]
     #[should_panic(expected = r#"Can be performed only by admin"#)]
@@ -86,4 +106,33 @@ mod test {
         let context = Context::new(admin);
         context.contract().update_contract(vec![], None);
     }
+
+    #[test]
+    fn test_assert_gas() {
+        const GAS_FOR_ASSERT_CALL: u64 = 529536222;
+
+        expect_panic(
+            &(),
+            "Not enough gas left. Consider attaching more gas to the transaction.",
+            || {
+                assert_gas(u64::MAX, || "Error message");
+            },
+        );
+
+        let gas_left = env::prepaid_gas().as_gas() - env::used_gas().as_gas();
+        expect_panic(&(), &format!("Need additional {GAS_FOR_ASSERT_CALL} gas"), || {
+            assert_gas(gas_left, || "Error message");
+        });
+
+        let gas_left = env::prepaid_gas().as_gas() - env::used_gas().as_gas();
+        expect_panic(&(), "Need additional 1 gas", || {
+            assert_gas(gas_left - GAS_FOR_ASSERT_CALL + 1, || "Error message");
+        });
+
+        let gas_left = env::prepaid_gas().as_gas() - env::used_gas().as_gas();
+        assert_gas(gas_left - GAS_FOR_ASSERT_CALL, || "Error message");
+
+        let gas_left = env::prepaid_gas().as_gas() - env::used_gas().as_gas();
+        assert_gas(gas_left - GAS_FOR_ASSERT_CALL - 1, || "Error message");
+    }
 }

contract/src/jar/api.rs

@@ -14,7 +14,7 @@ use crate::{
 };
 
 impl Contract {
-    fn can_be_restacked(&self, jar: &Jar, now: u64) -> bool {
+    fn can_be_restaked(&self, jar: &Jar, now: u64) -> bool {
         let product = self.get_product(&jar.product_id);
         !jar.is_empty() && product.is_enabled && product.allows_restaking() && jar.is_liquidable(&product, now)
     }
@@ -170,7 +170,7 @@ impl JarApi for Contract {
             })
             .jars
             .iter()
-            .filter(|j| self.can_be_restacked(j, now))
+            .filter(|j| self.can_be_restaked(j, now))
             .cloned()
             .collect();
 

contract/src/jar/tests/restake_all.rs

@@ -13,68 +13,67 @@ fn restake_all() {
     let alice = alice();
     let admin = admin();
 
-    let restackable_product = generate_product("restakable_product")
+    let restakable_product = generate_product("restakable_product")
         .with_allows_restaking(true)
         .lockup_term(MS_IN_YEAR);
 
-    let disabled_restackable_product = generate_product("disabled_restackable_product")
+    let disabled_restakable_product = generate_product("disabled_restakable_product")
         .with_allows_restaking(true)
         .enabled(false)
         .lockup_term(MS_IN_YEAR);
 
-    let non_restackable_product = generate_product("non_restakable_product")
+    let non_restakable_product = generate_product("non_restakable_product")
         .with_allows_restaking(false)
         .lockup_term(MS_IN_YEAR);
 
-    let long_term_restackable_product = generate_product("long_term_restackable_product")
+    let long_term_restakable_product = generate_product("long_term_restakable_product")
         .with_allows_restaking(true)
         .lockup_term(MS_IN_YEAR * 2);
 
-    let restackable_jar_1 = Jar::generate(JAR_ID_RANGE.fake(), &alice, &restackable_product.id).principal(PRINCIPAL);
-    let restackable_jar_2 = Jar::generate(JAR_ID_RANGE.fake(), &alice, &restackable_product.id).principal(PRINCIPAL);
+    let restakable_jar_1 = Jar::generate(JAR_ID_RANGE.fake(), &alice, &restakable_product.id).principal(PRINCIPAL);
+    let restakable_jar_2 = Jar::generate(JAR_ID_RANGE.fake(), &alice, &restakable_product.id).principal(PRINCIPAL);
 
-    let disabled_jar =
-        Jar::generate(JAR_ID_RANGE.fake(), &alice, &disabled_restackable_product.id).principal(PRINCIPAL);
+    let disabled_jar = Jar::generate(JAR_ID_RANGE.fake(), &alice, &disabled_restakable_product.id).principal(PRINCIPAL);
 
-    let non_restackable_jar =
-        Jar::generate(JAR_ID_RANGE.fake(), &alice, &non_restackable_product.id).principal(PRINCIPAL);
+    let non_restakable_jar =
+        Jar::generate(JAR_ID_RANGE.fake(), &alice, &non_restakable_product.id).principal(PRINCIPAL);
 
     let long_term_jar =
-        Jar::generate(JAR_ID_RANGE.fake(), &alice, &long_term_restackable_product.id).principal(PRINCIPAL);
+        Jar::generate(JAR_ID_RANGE.fake(), &alice, &long_term_restakable_product.id).principal(PRINCIPAL);
 
     let mut context = Context::new(admin)
         .with_products(&[
-            restackable_product,
-            disabled_restackable_product,
-            non_restackable_product,
-            long_term_restackable_product,
+            restakable_product,
+            disabled_restakable_product,
+            non_restakable_product,
+            long_term_restakable_product,
         ])
         .with_jars(&[
-            restackable_jar_1.clone(),
-            restackable_jar_2.clone(),
+            restakable_jar_1.clone(),
+            restakable_jar_2.clone(),
             disabled_jar.clone(),
-            non_restackable_jar.clone(),
+            non_restakable_jar.clone(),
             long_term_jar.clone(),
         ]);
 
     context.set_block_timestamp_in_days(366);
 
     context.switch_account(&alice);
 
-    let restacked_jars = context.contract().restake_all();
+    let restaked_jars = context.contract().restake_all();
 
-    assert_eq!(restacked_jars.len(), 2);
-    assert_eq!(restacked_jars.iter().map(|j| j.id.0).collect::<Vec<_>>(), vec![1, 2]);
+    assert_eq!(restaked_jars.len(), 2);
+    assert_eq!(restaked_jars.iter().map(|j| j.id.0).collect::<Vec<_>>(), vec![1, 2]);
 
     let all_jars = context.contract().get_jars_for_account(alice);
 
     assert_eq!(
         all_jars.iter().map(|j| j.id.0).collect::<Vec<_>>(),
         [
-            restackable_jar_1.id,
-            restackable_jar_2.id,
+            restakable_jar_1.id,
+            restakable_jar_2.id,
             disabled_jar.id,
-            non_restackable_jar.id,
+            non_restakable_jar.id,
             long_term_jar.id,
             1,
             2,

contract/src/test_utils.rs

@@ -79,6 +79,10 @@ pub trait AfterCatchUnwind {
     fn after_catch_unwind(&self);
 }
 
+impl AfterCatchUnwind for () {
+    fn after_catch_unwind(&self) {}
+}
+
 pub fn expect_panic(ctx: &impl AfterCatchUnwind, msg: &str, action: impl FnOnce() + UnwindSafe) {
     let res = catch_unwind(move || action());
 

contract/src/withdraw/api.rs

@@ -1,5 +1,4 @@
 use near_sdk::{
-    env::panic_str,
     ext_contract, is_promise_success,
     json_types::U128,
     near_bindgen,
@@ -13,7 +12,7 @@ use sweat_jar_model::{
     TokenAmount,
 };
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Debug)]
 #[serde(crate = "near_sdk::serde")]
 pub struct JarWithdraw {
     pub jar: Jar,
@@ -86,12 +85,11 @@ impl WithdrawApi for Contract {
         self.migrate_account_jars_if_needed(account_id.clone());
         let now = env::block_timestamp_ms();
 
-        let jars: Vec<JarWithdraw> = self
-            .account_jars
-            .get(&account_id)
-            .unwrap_or_else(|| {
-                panic_str(&format!("Jars for account '{account_id}' don't exist"));
-            })
+        let Some(account_jars) = self.account_jars.get(&account_id) else {
+            return PromiseOrValue::Value(BulkWithdrawView::default());
+        };
+
+        let jars: Vec<JarWithdraw> = account_jars
             .jars
             .clone()
             .into_iter()
@@ -104,6 +102,10 @@ impl WithdrawApi for Contract {
 
                 let amount = jar.principal;
 
+                if amount == 0 {
+                    return None;
+                }
+
                 let mut withdrawn_jar = jar.withdrawn(&product, amount, now);
                 let should_be_closed = withdrawn_jar.should_be_closed(&product, now);
 
@@ -120,6 +122,10 @@ impl WithdrawApi for Contract {
             })
             .collect();
 
+        if jars.is_empty() {
+            return PromiseOrValue::Value(BulkWithdrawView::default());
+        }
+
         self.transfer_bulk_withdraw(&account_id, jars)
     }
 }
@@ -270,14 +276,11 @@ impl Contract {
 
         let total_amount = jars.iter().map(|j| j.amount).sum();
 
-        let gas_left = crate::env::prepaid_gas().as_gas() - crate::env::used_gas().as_gas();
-
-        if gas_left
-            < crate::common::gas_data::GAS_FOR_FT_TRANSFER.as_gas()
-                + crate::common::gas_data::GAS_FOR_BULK_AFTER_WITHDRAW.as_gas()
-        {
-            panic_str("Not enough gas left to complete transfer_bulk_withdraw.");
-        }
+        crate::internal::assert_gas(
+            crate::common::gas_data::GAS_FOR_FT_TRANSFER.as_gas()
+                + crate::common::gas_data::GAS_FOR_BULK_AFTER_WITHDRAW.as_gas(),
+            || format!("transfer_bulk_withdraw. Number of jars: {}", jars.len()),
+        );
 
         self.ft_contract()
             .ft_transfer(account_id, total_amount, "bulk_withdraw", &total_fee)
@@ -344,6 +347,7 @@ impl Contract {
 }
 
 #[near_bindgen]
+#[mutants::skip] // Covered by integration tests
 impl WithdrawCallbacks for Contract {
     #[private]
     fn after_withdraw(