@@ -1825,3 +1825,236 @@ fn test_claim_root_default_mode_keep() {
18251825 assert_eq!(RootClaimType::<Test>::get(coldkey), RootClaimTypeEnum::Swap);
18261826 });
18271827}
1828+
1829+ #[test]
1830+ fn test_claim_root_with_moved_stake() {
1831+ new_test_ext(1).execute_with(|| {
1832+ let owner_coldkey = U256::from(1001);
1833+ let hotkey = U256::from(1002);
1834+ let alice_coldkey = U256::from(1003);
1835+ let bob_coldkey = U256::from(1004);
1836+ let eve_coldkey = U256::from(1005);
1837+ let netuid = add_dynamic_network(&hotkey, &owner_coldkey);
1838+
1839+ SubtensorModule::set_tao_weight(u64::MAX); // Set TAO weight to 1.0
1840+ SubtokenEnabled::<Test>::insert(NetUid::ROOT, true);
1841+ NetworksAdded::<Test>::insert(NetUid::ROOT, true);
1842+
1843+ let root_stake = 8_000_000u64;
1844+ SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
1845+ &hotkey,
1846+ &alice_coldkey,
1847+ NetUid::ROOT,
1848+ root_stake.into(),
1849+ );
1850+ SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
1851+ &hotkey,
1852+ &bob_coldkey,
1853+ NetUid::ROOT,
1854+ root_stake.into(),
1855+ );
1856+
1857+ let initial_total_hotkey_alpha = 10_000_000u64;
1858+ SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
1859+ &hotkey,
1860+ &owner_coldkey,
1861+ netuid,
1862+ initial_total_hotkey_alpha.into(),
1863+ );
1864+
1865+ // Claim root alpha
1866+
1867+ assert_ok!(SubtensorModule::set_root_claim_type(
1868+ RuntimeOrigin::signed(alice_coldkey),
1869+ RootClaimTypeEnum::Keep
1870+ ),);
1871+ assert_ok!(SubtensorModule::set_root_claim_type(
1872+ RuntimeOrigin::signed(bob_coldkey),
1873+ RootClaimTypeEnum::Keep
1874+ ),);
1875+
1876+ assert_ok!(SubtensorModule::set_root_claim_type(
1877+ RuntimeOrigin::signed(eve_coldkey),
1878+ RootClaimTypeEnum::Keep
1879+ ),);
1880+
1881+ // Distribute pending root alpha
1882+
1883+ let pending_root_alpha = 10_000_000u64;
1884+ SubtensorModule::distribute_emission(
1885+ netuid,
1886+ AlphaCurrency::ZERO,
1887+ AlphaCurrency::ZERO,
1888+ pending_root_alpha.into(),
1889+ AlphaCurrency::ZERO,
1890+ );
1891+
1892+ assert_ok!(SubtensorModule::claim_root(
1893+ RuntimeOrigin::signed(alice_coldkey),
1894+ BTreeSet::from([netuid])
1895+ ));
1896+ assert_ok!(SubtensorModule::claim_root(
1897+ RuntimeOrigin::signed(bob_coldkey),
1898+ BTreeSet::from([netuid])
1899+ ));
1900+
1901+ // Check stakes
1902+ let validator_take_percent = 0.18f64;
1903+
1904+ let alice_stake: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
1905+ &hotkey,
1906+ &alice_coldkey,
1907+ netuid,
1908+ )
1909+ .into();
1910+
1911+ let bob_stake: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
1912+ &hotkey,
1913+ &bob_coldkey,
1914+ netuid,
1915+ )
1916+ .into();
1917+
1918+ let estimated_stake = (pending_root_alpha as f64) * (1f64 - validator_take_percent) / 2f64;
1919+
1920+ assert_eq!(alice_stake, bob_stake);
1921+
1922+ assert_abs_diff_eq!(alice_stake, estimated_stake as u64, epsilon = 100u64,);
1923+
1924+
1925+ // Distribute pending root alpha
1926+
1927+ let pending_root_alpha = 10_000_000u64;
1928+ SubtensorModule::distribute_emission(
1929+ netuid,
1930+ AlphaCurrency::ZERO,
1931+ AlphaCurrency::ZERO,
1932+ pending_root_alpha.into(),
1933+ AlphaCurrency::ZERO,
1934+ );
1935+
1936+ // Transfer stake to other coldkey
1937+ let stake_decrement = root_stake / 2u64;
1938+
1939+ assert_ok!(SubtensorModule::transfer_stake(
1940+ RuntimeOrigin::signed(bob_coldkey,),
1941+ eve_coldkey,
1942+ hotkey,
1943+ NetUid::ROOT,
1944+ NetUid::ROOT,
1945+ stake_decrement.into(),
1946+ ));
1947+
1948+ let eve_stake: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
1949+ &hotkey,
1950+ &eve_coldkey,
1951+ netuid,
1952+ )
1953+ .into();
1954+
1955+ assert_ok!(SubtensorModule::claim_root(
1956+ RuntimeOrigin::signed(alice_coldkey),
1957+ BTreeSet::from([netuid])
1958+ ));
1959+ assert_ok!(SubtensorModule::claim_root(
1960+ RuntimeOrigin::signed(bob_coldkey),
1961+ BTreeSet::from([netuid])
1962+ ));
1963+
1964+ assert_ok!(SubtensorModule::claim_root(
1965+ RuntimeOrigin::signed(eve_coldkey),
1966+ BTreeSet::from([netuid])
1967+ ));
1968+
1969+ // Check new stakes
1970+
1971+ let alice_stake2: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
1972+ &hotkey,
1973+ &alice_coldkey,
1974+ netuid,
1975+ )
1976+ .into();
1977+
1978+ let bob_stake2: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
1979+ &hotkey,
1980+ &bob_coldkey,
1981+ netuid,
1982+ )
1983+ .into();
1984+
1985+ let eve_stake2: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
1986+ &hotkey,
1987+ &eve_coldkey,
1988+ netuid,
1989+ )
1990+ .into();
1991+
1992+ // Eve should not have gotten any root claim
1993+ let eve_stake_diff = eve_stake2 - eve_stake;
1994+ assert_abs_diff_eq!(eve_stake_diff, 0, epsilon = 100u64,);
1995+
1996+ let estimated_stake = (pending_root_alpha as f64) * (1f64 - validator_take_percent) / 2f64;
1997+
1998+ let alice_stake_diff = alice_stake2 - alice_stake;
1999+ let bob_stake_diff = bob_stake2 - bob_stake;
2000+
2001+ assert_abs_diff_eq!(alice_stake_diff, bob_stake_diff, epsilon = 100u64,);
2002+ assert_abs_diff_eq!(bob_stake_diff, estimated_stake as u64, epsilon = 100u64,);
2003+
2004+ // Transfer stake back
2005+ let stake_increment = stake_decrement;
2006+
2007+ assert_ok!(SubtensorModule::transfer_stake(
2008+ RuntimeOrigin::signed(eve_coldkey,),
2009+ bob_coldkey,
2010+ hotkey,
2011+ NetUid::ROOT,
2012+ NetUid::ROOT,
2013+ stake_increment.into(),
2014+ ));
2015+
2016+ // Distribute pending root alpha
2017+
2018+ let pending_root_alpha = 10_000_000u64;
2019+ SubtensorModule::distribute_emission(
2020+ netuid,
2021+ AlphaCurrency::ZERO,
2022+ AlphaCurrency::ZERO,
2023+ pending_root_alpha.into(),
2024+ AlphaCurrency::ZERO,
2025+ );
2026+
2027+ assert_ok!(SubtensorModule::claim_root(
2028+ RuntimeOrigin::signed(alice_coldkey),
2029+ BTreeSet::from([netuid])
2030+ ));
2031+ assert_ok!(SubtensorModule::claim_root(
2032+ RuntimeOrigin::signed(bob_coldkey),
2033+ BTreeSet::from([netuid])
2034+ ));
2035+
2036+ // Check new stakes
2037+
2038+ let alice_stake3: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
2039+ &hotkey,
2040+ &alice_coldkey,
2041+ netuid,
2042+ )
2043+ .into();
2044+
2045+ let bob_stake3: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
2046+ &hotkey,
2047+ &bob_coldkey,
2048+ netuid,
2049+ )
2050+ .into();
2051+
2052+ let estimated_stake = (pending_root_alpha as f64) * (1f64 - validator_take_percent) / 2f64;
2053+
2054+ let alice_stake_diff2 = alice_stake3 - alice_stake2;
2055+ let bob_stake_diff2 = bob_stake3 - bob_stake2;
2056+
2057+ assert_abs_diff_eq!(alice_stake_diff2, bob_stake_diff2, epsilon = 100u64,);
2058+ assert_abs_diff_eq!(bob_stake_diff2, estimated_stake as u64, epsilon = 100u64,);
2059+ });
2060+ }
0 commit comments