1
// Copyright 2019-2025 PureStake Inc.
2
// This file is part of Moonbeam.
3

            
4
// Moonbeam is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Moonbeam is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
16

            
17
//! Test primitives for Moonbeam, including mock implementations for testing.
18

            
19
use frame_support::weights::{constants::WEIGHT_REF_TIME_PER_SECOND, Weight};
20
use sp_runtime::DispatchError;
21
use sp_std::cell::RefCell;
22
use sp_std::collections::btree_map::BTreeMap;
23
use xcm::latest::Location;
24
use xcm_primitives::XcmFeeTrader;
25

            
26
thread_local! {
27
	/// Thread-local storage for fee-per-second values keyed by asset location.
28
	static FEE_PER_SECOND: RefCell<BTreeMap<Location, u128>> = RefCell::new(BTreeMap::new());
29
}
30

            
31
/// Memory-based fee trader for tests that stores fee-per-second values in memory.
32
///
33
/// This implementation calculates fees based on weight and a fee-per-second rate,
34
/// similar to how the real weight-trader pallet works. It's useful for testing
35
/// XCM fee payment scenarios without requiring a full runtime setup.
36
pub struct MemoryFeeTrader;
37

            
38
impl XcmFeeTrader for MemoryFeeTrader {
39
273
	fn compute_fee(
40
273
		weight: Weight,
41
273
		asset_location: &Location,
42
273
		explicit_amount: Option<u128>,
43
273
	) -> Result<u128, DispatchError> {
44
		// If explicit amount is provided, use it directly
45
273
		if let Some(amount) = explicit_amount {
46
			return Ok(amount);
47
273
		}
48

            
49
		// Get fee-per-second from storage
50
273
		let fee_per_second = FEE_PER_SECOND
51
273
			.with(|map| map.borrow().get(asset_location).copied())
52
273
			.ok_or(DispatchError::Other("Fee per second not set"))?;
53

            
54
		// Calculate fee using the same formula as the weight-trader pallet
55
		// fee = (fee_per_second * weight.ref_time() + weight_per_second - 1) / weight_per_second
56
252
		let weight_per_second_u128 = WEIGHT_REF_TIME_PER_SECOND as u128;
57
252
		let fee_mul_rounded_up = (fee_per_second.saturating_mul(weight.ref_time() as u128))
58
252
			.saturating_add(weight_per_second_u128 - 1);
59
252
		Ok(fee_mul_rounded_up / weight_per_second_u128)
60
273
	}
61

            
62
21
	fn get_asset_price(asset_location: &Location) -> Option<u128> {
63
21
		FEE_PER_SECOND.with(|map| map.borrow().get(asset_location).copied())
64
21
	}
65

            
66
315
	fn set_asset_price(asset_location: Location, value: u128) -> Result<(), DispatchError> {
67
315
		FEE_PER_SECOND.with(|map| {
68
315
			map.borrow_mut().insert(asset_location, value);
69
315
		});
70
315
		Ok(())
71
315
	}
72

            
73
	fn remove_asset(asset_location: Location) -> Result<(), DispatchError> {
74
		FEE_PER_SECOND.with(|map| {
75
			map.borrow_mut().remove(&asset_location);
76
		});
77
		Ok(())
78
	}
79
}