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
83
84
85
86
87
88
89

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

use core::objects::fee::{Fee, UsageFee};
use cosmwasm_std::{CosmosMsg, Uint128};
use cw_asset::Asset;

use crate::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 super::*;
    use cosmwasm_std::{testing::MockApi, Addr, Decimal};
    use cw_asset::AssetInfo;

    // 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(
            &MockApi::default(),
            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(
            &MockApi::default(),
            Decimal::zero(),
            Addr::unchecked("recipient"),
        )
        .unwrap();

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