1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

//! # Fee helpers
//! Helper trait that lets you easily charge fees on assets

use cosmwasm_std::{CosmosMsg, Uint128};
use cw_asset::Asset;

use crate::{
    std::objects::fee::{Fee, UsageFee},
    AbstractSdkResult,
};

/// Indicates that the implementing type can be charged fees.
pub trait Chargeable {
    /// Charge a fee on the asset and returns the amount charged.
    fn charge_fee(&mut self, fee: Fee) -> AbstractSdkResult<Uint128>;
    /// Charge a fee on the asset and returns the fee transfer message.
    fn charge_usage_fee(&mut self, fee: UsageFee) -> AbstractSdkResult<Option<CosmosMsg>>;
}

impl Chargeable for Asset {
    fn charge_fee(&mut self, fee: Fee) -> AbstractSdkResult<Uint128> {
        let fee_amount = fee.compute(self.amount);
        self.amount -= fee_amount;
        Ok(fee_amount)
    }

    /// returns a fee message if fee > 0
    fn charge_usage_fee(&mut self, fee: UsageFee) -> AbstractSdkResult<Option<CosmosMsg>> {
        let fee_amount = fee.compute(self.amount);
        if fee_amount.is_zero() {
            return Ok(None);
        }
        self.amount -= fee_amount;
        Ok(Some(
            Asset::new(self.info.clone(), fee_amount).transfer_msg(fee.recipient())?,
        ))
    }
}

#[cfg(test)]
mod tests {
    use cosmwasm_std::{Addr, Decimal};
    use cw_asset::AssetInfo;

    use super::*;

    // test that we can charge fees on assets

    #[test]
    fn test_charge_fee() {
        let info = AssetInfo::native("uusd");
        let mut asset = Asset::new(info, 1000u128);
        let fee = Fee::new(Decimal::percent(10)).unwrap();
        let charged = asset.charge_fee(fee).unwrap();
        assert_eq!(asset.amount.u128(), 900);
        assert_eq!(charged.u128(), 100);
    }
    // test transfer fee
    #[test]
    fn test_charge_transfer_fee() {
        let info = AssetInfo::native("uusd");
        let mut asset: Asset = Asset::new(info.clone(), 1000u128);
        let fee = UsageFee::new(Decimal::percent(10), Addr::unchecked("recipient")).unwrap();
        let msg = asset.charge_usage_fee(fee).unwrap();
        assert_eq!(asset.amount.u128(), 900);
        assert_eq!(
            msg,
            Some(
                Asset::new(info, 100u128)
                    .transfer_msg(Addr::unchecked("recipient"))
                    .unwrap()
            )
        );

        // test zero fee
        let fee = UsageFee::new(Decimal::zero(), Addr::unchecked("recipient")).unwrap();

        let msg = asset.charge_usage_fee(fee).unwrap();
        assert_eq!(asset.amount.u128(), 900);
        assert_eq!(msg, None);
    }
}