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

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

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

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<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)
    }

    fn charge_usage_fee(&mut self, fee: UsageFee) -> AbstractSdkResult<CosmosMsg> {
        let fee_amount = fee.compute(self.amount);
        self.amount -= fee_amount;
        Ok(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,
            Asset::new(info, 100u128)
                .transfer_msg(Addr::unchecked("recipient"))
                .unwrap()
        );
    }
}