Feature claim and withdraw

ref-farming/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ref_farming"
-version = "1.0.1"
+version = "1.0.2"
 authors = ["Marco Sun <[email protected]>"]
 edition = "2018"
 

ref-farming/src/actions_of_reward.rs

@@ -1,7 +1,7 @@
-
-use std::convert::TryInto;
+use std::collections::HashMap;
 use near_sdk::json_types::{ValidAccountId, U128};
-use near_sdk::{assert_one_yocto, env, near_bindgen, AccountId, Balance, PromiseResult};
+use near_sdk::{assert_one_yocto, env, near_bindgen, 
+    AccountId, Balance, Promise, PromiseResult};
 
 use crate::utils::{ext_fungible_token, ext_self, GAS_FOR_FT_TRANSFER, GAS_FOR_RESOLVE_TRANSFER, parse_farm_id};
 use crate::errors::*;
@@ -44,9 +44,30 @@ impl Contract {
         self.assert_storage_usage(&sender_id);
     }
 
+    /// if withdraw_all_tokens is true, withdraw all tokens, 
+    /// or just withdraw the reward token of this farm.
+    pub fn claim_and_withdraw_by_farm(&mut self, farm_id: FarmId, withdraw_all_tokens: bool) {
+        let sender_id = env::predecessor_account_id();
+        self.claim_reward_by_farm(farm_id.clone());
+        if withdraw_all_tokens {
+            self.internal_withdraw_all_rewards(&sender_id, None);
+        } else {
+            if let Some(farm) = self.data().farms.get(&farm_id) {
+                let token_id = farm.get_reward_token();
+                self.internal_withdraw_all_rewards(&sender_id, Some(&token_id));
+            }
+        }
+    }
+
+    pub fn claim_and_withdraw_by_seed(&mut self, seed_id: SeedId) {
+        let sender_id = env::predecessor_account_id();
+        self.claim_reward_by_seed(seed_id.clone());
+        self.internal_withdraw_all_rewards(&sender_id, None);
+    }
+
     /// Withdraws given reward token of given user.
     #[payable]
-    pub fn withdraw_reward(&mut self, token_id: ValidAccountId, amount: Option<U128>) {
+    pub fn withdraw_reward(&mut self, token_id: ValidAccountId, amount: Option<U128>) -> Promise {
         assert_one_yocto();
 
         let token_id: AccountId = token_id.into();
@@ -59,22 +80,9 @@ impl Contract {
         // Note: subtraction, will be reverted if the promise fails.
         let amount = farmer.get_ref_mut().sub_reward(&token_id, amount);
         self.data_mut().farmers.insert(&sender_id, &farmer);
-        ext_fungible_token::ft_transfer(
-            sender_id.clone().try_into().unwrap(),
-            amount.into(),
-            None,
-            &token_id,
-            1,
-            GAS_FOR_FT_TRANSFER,
-        )
-        .then(ext_self::callback_post_withdraw_reward(
-            token_id,
-            sender_id,
-            amount.into(),
-            &env::current_account_id(),
-            0,
-            GAS_FOR_RESOLVE_TRANSFER,
-        ));
+
+        self.internal_send_tokens(&sender_id, &token_id, amount)
+
     }
 
     #[private]
@@ -153,6 +161,59 @@ fn claim_user_reward_from_farm(
 }
 
 impl Contract {
+
+    /// if token was given, then just withdraw that token, or withdraw all tokens
+    pub(crate) fn internal_withdraw_all_rewards(&mut self, farmer_id: &AccountId, token: Option<&AccountId>) {
+
+        let view_farmer = self.get_farmer(&farmer_id);
+        let tokens = {
+            if let Some(token) = token {
+                vec![token]
+            } else {
+                view_farmer.get_ref().rewards.keys().collect::<Vec<_>>()
+            }
+        };
+
+        let mut farmer = self.get_farmer(&farmer_id);
+        let mut withdraw_tasks = HashMap::<AccountId, u128>::new();
+        for token_id in tokens {
+            let amount = farmer.get_ref_mut().sub_reward(&token_id, 0);
+            withdraw_tasks.insert(token_id.clone(), amount);
+        }
+        self.data_mut().farmers.insert(farmer_id, &farmer);
+
+        for (token_id, amount) in withdraw_tasks {
+            if amount > 0 {
+                self.internal_send_tokens(farmer_id, &token_id, amount);
+            }
+        }
+    }
+
+    /// Sends given amount to given user and if it fails, returns it back to user's balance.
+    /// Tokens must already be subtracted from internal balance.
+    pub(crate) fn internal_send_tokens(
+        &self,
+        sender_id: &AccountId,
+        token_id: &AccountId,
+        amount: Balance,
+    ) -> Promise {
+        ext_fungible_token::ft_transfer(
+            sender_id.clone(),
+            U128(amount),
+            None,
+            token_id,
+            1,
+            GAS_FOR_FT_TRANSFER,
+        )
+        .then(ext_self::callback_post_withdraw_reward(
+            token_id.clone(),
+            sender_id.clone(),
+            U128(amount),
+            &env::current_account_id(),
+            0,
+            GAS_FOR_RESOLVE_TRANSFER,
+        ))
+    }
 
     pub(crate) fn internal_claim_user_reward_by_seed_id(
         &mut self, 

ref-farming/src/farmer.rs

@@ -52,7 +52,7 @@ impl Farmer {
     pub(crate) fn sub_reward(&mut self, token: &AccountId, amount: Balance) -> Balance {
         let value = *self.rewards.get(token).expect(ERR21_TOKEN_NOT_REG);
         assert!(value >= amount, "{}", ERR22_NOT_ENOUGH_TOKENS);
-        if amount == 0 {
+        if amount == 0 || amount == value {
             self.rewards.remove(&token.clone());
             value
         } else {

ref-farming/tests/common/views.rs

@@ -4,6 +4,7 @@ use near_sdk_sim::{view, ContractAccount};
 
 use super::utils::to_va;
 use ref_farming::{ContractContract as Farming, FarmInfo};
+use test_token::ContractContract as TestToken;
 use std::collections::HashMap;
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
@@ -184,6 +185,12 @@ pub(crate) fn show_storage_balance(farming: &ContractAccount<Farming>, farmer: S
     ret
 }
 
+pub(crate) fn balance_of(token: &ContractAccount<TestToken>, account_id: String) -> u128 {
+    view!(token.ft_balance_of(to_va(account_id.clone())))
+        .unwrap_json::<U128>()
+        .0
+}
+
 // =============  Assertions  ===============
 #[allow(dead_code)]
 pub(crate) fn assert_farming(