Skip to content

Commit 8f197a2

Browse files
camfairchildcamfairchild
committed
add tests for validation filter
1 parent e538804 commit 8f197a2

File tree

1 file changed

+326
-1
lines changed

1 file changed

+326
-1
lines changed

‎pallets/subtensor/src/tests/swap_coldkey.rs‎

Lines changed: 326 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1726,7 +1726,7 @@ fn test_schedule_swap_coldkey_duplicate() {
17261726
});
17271727
}
17281728

1729-
// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap_coldkey -- test_schedule_swap_coldkey_execution --exact --nocapture
1729+
// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --package pallet-subtensor --lib -- tests::swap_coldkey::test_schedule_swap_coldkey_execution --exact --show-output --nocapture
17301730
#[test]
17311731
fn test_schedule_swap_coldkey_execution() {
17321732
new_test_ext(1).execute_with(|| {
@@ -2022,3 +2022,328 @@ fn test_cant_schedule_swap_without_enough_to_burn() {
20222022
);
20232023
});
20242024
}
2025+
2026+
// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --package pallet-subtensor --lib -- tests::swap_coldkey::test_coldkey_in_swap_schedule_prevents_funds_usage --exact --show-output --nocapture
2027+
#[test]
2028+
fn test_coldkey_in_swap_schedule_prevents_funds_usage() {
2029+
// Testing the signed extension validate function
2030+
// correctly filters transactions that attempt to use funds
2031+
// while a coldkey swap is scheduled.
2032+
2033+
new_test_ext(0).execute_with(|| {
2034+
let netuid: u16 = 1;
2035+
let version_key: u64 = 0;
2036+
let coldkey = U256::from(0);
2037+
let new_coldkey = U256::from(1);
2038+
let hotkey: U256 = U256::from(2); // Add the hotkey field
2039+
assert_ne!(hotkey, coldkey); // Ensure hotkey is NOT the same as coldkey !!!
2040+
let fee = DefaultStakingFee::<Test>::get();
2041+
2042+
let who = coldkey; // The coldkey signs this transaction
2043+
2044+
// Disallowed transactions are
2045+
// - add_stake
2046+
// - add_stake_limit
2047+
// - swap_stake
2048+
// - swap_stake_limit
2049+
// - move_stake
2050+
// - transfer_stake
2051+
// - balances.transfer_all
2052+
// - balances.transfer_allow_death
2053+
// - balances.transfer_keep_alive
2054+
2055+
// Allowed transactions are:
2056+
// - remove_stake
2057+
// - remove_stake_limit
2058+
// others...
2059+
2060+
// Create netuid
2061+
add_network(netuid, 1, 0);
2062+
// Register the hotkey
2063+
SubtensorModule::append_neuron(netuid, &hotkey, 0);
2064+
crate::Owner::<Test>::insert(hotkey, coldkey);
2065+
2066+
SubtensorModule::add_balance_to_coldkey_account(&who, u64::MAX);
2067+
2068+
// Set the minimum stake to 0.
2069+
SubtensorModule::set_stake_threshold(0);
2070+
// Add stake to the hotkey
2071+
assert_ok!(SubtensorModule::add_stake(
2072+
<<Test as Config>::RuntimeOrigin>::signed(who),
2073+
hotkey,
2074+
netuid,
2075+
100_000_000_000
2076+
));
2077+
2078+
// Schedule the coldkey for a swap
2079+
assert_ok!(SubtensorModule::schedule_swap_coldkey(
2080+
<<Test as Config>::RuntimeOrigin>::signed(who),
2081+
new_coldkey
2082+
));
2083+
2084+
assert!(ColdkeySwapScheduled::<Test>::contains_key(who));
2085+
2086+
// Setup the extension
2087+
let info: crate::DispatchInfo =
2088+
crate::DispatchInfoOf::<<Test as frame_system::Config>::RuntimeCall>::default();
2089+
let extension = crate::SubtensorSignedExtension::<Test>::new();
2090+
2091+
// Try each call
2092+
2093+
// Add stake
2094+
let call = RuntimeCall::SubtensorModule(SubtensorCall::add_stake {
2095+
hotkey,
2096+
netuid,
2097+
amount_staked: 100_000_000_000,
2098+
});
2099+
let result: Result<ValidTransaction, TransactionValidityError> =
2100+
extension.validate(&who, &call.clone(), &info, 10);
2101+
// Should fail
2102+
assert_err!(
2103+
// Should get an invalid transaction error
2104+
result,
2105+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2106+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2107+
))
2108+
);
2109+
2110+
// Add stake limit
2111+
let call = RuntimeCall::SubtensorModule(SubtensorCall::add_stake_limit {
2112+
hotkey,
2113+
netuid,
2114+
amount_staked: 100_000_000_000,
2115+
limit_price: 100_000_000_000,
2116+
allow_partial: false,
2117+
});
2118+
let result = extension.validate(&who, &call.clone(), &info, 10);
2119+
// Should fail
2120+
assert_err!(
2121+
// Should get an invalid transaction error
2122+
result,
2123+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2124+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2125+
))
2126+
);
2127+
2128+
// Swap stake
2129+
let call = RuntimeCall::SubtensorModule(SubtensorCall::swap_stake {
2130+
hotkey,
2131+
origin_netuid: netuid,
2132+
destination_netuid: netuid,
2133+
alpha_amount: 100_000_000_000,
2134+
});
2135+
let result = extension.validate(&who, &call.clone(), &info, 10);
2136+
// Should fail
2137+
assert_err!(
2138+
// Should get an invalid transaction error
2139+
result,
2140+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2141+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2142+
))
2143+
);
2144+
2145+
// Swap stake limit
2146+
let call = RuntimeCall::SubtensorModule(SubtensorCall::swap_stake_limit {
2147+
hotkey,
2148+
origin_netuid: netuid,
2149+
destination_netuid: netuid,
2150+
alpha_amount: 100_000_000_000,
2151+
limit_price: 100_000_000_000,
2152+
allow_partial: false,
2153+
});
2154+
let result = extension.validate(&who, &call.clone(), &info, 10);
2155+
// Should fail
2156+
assert_err!(
2157+
// Should get an invalid transaction error
2158+
result,
2159+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2160+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2161+
))
2162+
);
2163+
2164+
// Move stake
2165+
let call = RuntimeCall::SubtensorModule(SubtensorCall::move_stake {
2166+
origin_hotkey: hotkey,
2167+
destination_hotkey: hotkey,
2168+
origin_netuid: netuid,
2169+
destination_netuid: netuid,
2170+
alpha_amount: 100_000_000_000,
2171+
});
2172+
let result = extension.validate(&who, &call.clone(), &info, 10);
2173+
// Should fail
2174+
assert_err!(
2175+
// Should get an invalid transaction error
2176+
result,
2177+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2178+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2179+
))
2180+
);
2181+
2182+
// Transfer stake
2183+
let call = RuntimeCall::SubtensorModule(SubtensorCall::transfer_stake {
2184+
destination_coldkey: new_coldkey,
2185+
hotkey,
2186+
origin_netuid: netuid,
2187+
destination_netuid: netuid,
2188+
alpha_amount: 100_000_000_000,
2189+
});
2190+
let result = extension.validate(&who, &call.clone(), &info, 10);
2191+
// Should fail
2192+
assert_err!(
2193+
// Should get an invalid transaction error
2194+
result,
2195+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2196+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2197+
))
2198+
);
2199+
2200+
// Transfer all
2201+
let call = RuntimeCall::Balances(BalancesCall::transfer_all {
2202+
dest: new_coldkey,
2203+
keep_alive: false,
2204+
});
2205+
let result = extension.validate(&who, &call.clone(), &info, 10);
2206+
// Should fail
2207+
assert_err!(
2208+
// Should get an invalid transaction error
2209+
result,
2210+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2211+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2212+
))
2213+
);
2214+
2215+
// Transfer keep alive
2216+
let call = RuntimeCall::Balances(BalancesCall::transfer_keep_alive {
2217+
dest: new_coldkey,
2218+
value: 100_000_000_000,
2219+
});
2220+
let result = extension.validate(&who, &call.clone(), &info, 10);
2221+
// Should fail
2222+
assert_err!(
2223+
// Should get an invalid transaction error
2224+
result,
2225+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2226+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2227+
))
2228+
);
2229+
2230+
// Transfer allow death
2231+
let call = RuntimeCall::Balances(BalancesCall::transfer_allow_death {
2232+
dest: new_coldkey,
2233+
value: 100_000_000_000,
2234+
});
2235+
let result = extension.validate(&who, &call.clone(), &info, 10);
2236+
// Should fail
2237+
assert_err!(
2238+
// Should get an invalid transaction error
2239+
result,
2240+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2241+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2242+
))
2243+
);
2244+
2245+
// Burned register
2246+
let call = RuntimeCall::SubtensorModule(SubtensorCall::burned_register { netuid, hotkey });
2247+
let result = extension.validate(&who, &call.clone(), &info, 10);
2248+
// Should fail
2249+
assert_err!(
2250+
// Should get an invalid transaction error
2251+
result,
2252+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2253+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2254+
))
2255+
);
2256+
2257+
// Remove stake
2258+
let call = RuntimeCall::SubtensorModule(SubtensorCall::remove_stake {
2259+
hotkey,
2260+
netuid,
2261+
amount_unstaked: 1_000_000,
2262+
});
2263+
let result = extension.validate(&who, &call.clone(), &info, 10);
2264+
// Should pass, not in list.
2265+
assert_ok!(result);
2266+
2267+
// Remove stake limit
2268+
let call = RuntimeCall::SubtensorModule(SubtensorCall::remove_stake_limit {
2269+
hotkey,
2270+
netuid,
2271+
amount_unstaked: 1_000_000,
2272+
limit_price: 123456789, // should be low enough
2273+
allow_partial: true,
2274+
});
2275+
let result = extension.validate(&who, &call.clone(), &info, 10);
2276+
// Should pass, not in list.
2277+
assert_ok!(result);
2278+
});
2279+
}
2280+
2281+
// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --package pallet-subtensor --lib -- tests::swap_coldkey::test_coldkey_in_swap_schedule_prevents_critical_calls --exact --show-output --nocapture
2282+
#[test]
2283+
fn test_coldkey_in_swap_schedule_prevents_critical_calls() {
2284+
// Testing the signed extension validate function
2285+
// correctly filters transactions that are critical
2286+
// while a coldkey swap is scheduled.
2287+
2288+
new_test_ext(0).execute_with(|| {
2289+
let netuid: u16 = 1;
2290+
let version_key: u64 = 0;
2291+
let coldkey = U256::from(0);
2292+
let new_coldkey = U256::from(1);
2293+
let hotkey: U256 = U256::from(2); // Add the hotkey field
2294+
assert_ne!(hotkey, coldkey); // Ensure hotkey is NOT the same as coldkey !!!
2295+
let fee = DefaultStakingFee::<Test>::get();
2296+
2297+
let who = coldkey; // The coldkey signs this transaction
2298+
2299+
// Disallowed transactions are
2300+
// - dissolve_network
2301+
2302+
// Create netuid
2303+
add_network(netuid, 1, 0);
2304+
// Register the hotkey
2305+
SubtensorModule::append_neuron(netuid, &hotkey, 0);
2306+
crate::Owner::<Test>::insert(hotkey, coldkey);
2307+
2308+
SubtensorModule::add_balance_to_coldkey_account(&who, u64::MAX);
2309+
2310+
// Set the minimum stake to 0.
2311+
SubtensorModule::set_stake_threshold(0);
2312+
// Add stake to the hotkey
2313+
assert_ok!(SubtensorModule::add_stake(
2314+
<<Test as Config>::RuntimeOrigin>::signed(who),
2315+
hotkey,
2316+
netuid,
2317+
100_000_000_000
2318+
));
2319+
2320+
// Schedule the coldkey for a swap
2321+
assert_ok!(SubtensorModule::schedule_swap_coldkey(
2322+
<<Test as Config>::RuntimeOrigin>::signed(who),
2323+
new_coldkey
2324+
));
2325+
2326+
assert!(ColdkeySwapScheduled::<Test>::contains_key(who));
2327+
2328+
// Setup the extension
2329+
let info: crate::DispatchInfo =
2330+
crate::DispatchInfoOf::<<Test as frame_system::Config>::RuntimeCall>::default();
2331+
let extension = crate::SubtensorSignedExtension::<Test>::new();
2332+
2333+
// Try each call
2334+
2335+
// Dissolve network
2336+
let call =
2337+
RuntimeCall::SubtensorModule(SubtensorCall::dissolve_network { netuid, coldkey });
2338+
let result: Result<ValidTransaction, TransactionValidityError> =
2339+
extension.validate(&who, &call.clone(), &info, 10);
2340+
// Should fail
2341+
assert_err!(
2342+
// Should get an invalid transaction error
2343+
result,
2344+
crate::TransactionValidityError::Invalid(crate::InvalidTransaction::Custom(
2345+
CustomTransactionError::ColdkeyInSwapSchedule.into()
2346+
))
2347+
);
2348+
});
2349+
}

0 commit comments

Comments
 (0)