Merge e791c659af into dc1b8dfccd

Cargo.lock

@@ -8159,8 +8159,13 @@ version = "0.1.0"
 dependencies = [
  "frame-support",
  "frame-system",
+ "pallet-babe",
+ "pallet-grandpa",
+ "pallet-timestamp",
  "parity-scale-codec",
+ "rand_core",
  "scale-info",
+ "serai-abi",
  "serai-coins-pallet",
  "serai-dex-pallet",
  "serai-economic-security-pallet",
@@ -8169,6 +8174,8 @@ dependencies = [
  "serai-primitives",
  "serai-validator-sets-pallet",
  "serai-validator-sets-primitives",
+ "sp-core",
+ "sp-io",
  "sp-runtime",
  "sp-std",
 ]

substrate/client/tests/emissions.rs

@@ -1,251 +0,0 @@
-use std::{time::Duration, collections::HashMap};
-use rand_core::{RngCore, OsRng};
-
-use serai_client::TemporalSerai;
-
-use serai_abi::{
-  emissions::primitives::{INITIAL_REWARD_PER_BLOCK, SECURE_BY},
-  in_instructions::primitives::Batch,
-  primitives::{
-    BlockHash, ExternalBalance, ExternalCoin, ExternalNetworkId, EXTERNAL_NETWORKS,
-    FAST_EPOCH_DURATION, FAST_EPOCH_INITIAL_PERIOD, NETWORKS, TARGET_BLOCK_TIME, Amount, NetworkId,
-  },
-  validator_sets::primitives::Session,
-};
-
-use serai_client::Serai;
-
-mod common;
-use common::{genesis_liquidity::set_up_genesis, in_instructions::provide_batch};
-
-serai_test_fast_epoch!(
-  emissions: (|serai: Serai| async move {
-    test_emissions(serai).await;
-  })
-);
-
-async fn send_batches(serai: &Serai, ids: &mut HashMap<ExternalNetworkId, u32>) {
-  for network in EXTERNAL_NETWORKS {
-    // set up batch id
-    ids
-      .entry(network)
-      .and_modify(|v| {
-        *v += 1;
-      })
-      .or_insert(0);
-
-    // set up block hash
-    let mut block = BlockHash([0; 32]);
-    OsRng.fill_bytes(&mut block.0);
-
-    provide_batch(serai, Batch { network, id: ids[&network], block, instructions: vec![] }).await;
-  }
-}
-
-async fn test_emissions(serai: Serai) {
-  // set up the genesis
-  let values = HashMap::from([
-    (ExternalCoin::Monero, 184100),
-    (ExternalCoin::Ether, 4785000),
-    (ExternalCoin::Dai, 1500),
-  ]);
-  let (_, mut batch_ids) = set_up_genesis(&serai, &values).await;
-
-  // wait until genesis is complete
-  let mut genesis_complete_block = None;
-  while genesis_complete_block.is_none() {
-    tokio::time::sleep(Duration::from_secs(1)).await;
-    genesis_complete_block = serai
-      .as_of_latest_finalized_block()
-      .await
-      .unwrap()
-      .genesis_liquidity()
-      .genesis_complete_block()
-      .await
-      .unwrap();
-  }
-
-  for _ in 0 .. 3 {
-    // get current stakes
-    let mut current_stake = HashMap::new();
-    for n in NETWORKS {
-      // TODO: investigate why serai network TAS isn't visible at session 0.
-      let stake = serai
-        .as_of_latest_finalized_block()
-        .await
-        .unwrap()
-        .validator_sets()
-        .total_allocated_stake(n)
-        .await
-        .unwrap()
-        .unwrap_or(Amount(0))
-        .0;
-      current_stake.insert(n, stake);
-    }
-
-    // wait for a session change
-    let current_session = wait_for_session_change(&serai).await;
-
-    // get last block
-    let last_block = serai.latest_finalized_block().await.unwrap();
-    let serai_latest = serai.as_of(last_block.hash());
-    let change_block_number = last_block.number();
-
-    // get distances to ec security & block count of the previous session
-    let (distances, total_distance) = get_distances(&serai_latest, &current_stake).await;
-    let block_count = get_session_blocks(&serai_latest, current_session - 1).await;
-
-    // calculate how much reward in this session
-    let reward_this_epoch =
-      if change_block_number < (genesis_complete_block.unwrap() + FAST_EPOCH_INITIAL_PERIOD) {
-        block_count * INITIAL_REWARD_PER_BLOCK
-      } else {
-        let blocks_until = SECURE_BY - change_block_number;
-        let block_reward = total_distance / blocks_until;
-        block_count * block_reward
-      };
-
-    let reward_per_network = distances
-      .into_iter()
-      .map(|(n, distance)| {
-        let reward = u64::try_from(
-          u128::from(reward_this_epoch).saturating_mul(u128::from(distance)) /
-            u128::from(total_distance),
-        )
-        .unwrap();
-        (n, reward)
-      })
-      .collect::<HashMap<NetworkId, u64>>();
-
-    // retire the prev-set so that TotalAllocatedStake updated.
-    send_batches(&serai, &mut batch_ids).await;
-
-    for (n, reward) in reward_per_network {
-      let stake = serai
-        .as_of_latest_finalized_block()
-        .await
-        .unwrap()
-        .validator_sets()
-        .total_allocated_stake(n)
-        .await
-        .unwrap()
-        .unwrap_or(Amount(0))
-        .0;
-
-      // all reward should automatically staked for the network since we are in initial period.
-      assert_eq!(stake, *current_stake.get(&n).unwrap() + reward);
-    }
-
-    // TODO: check stake per address?
-    // TODO: check post ec security era
-  }
-}
-
-/// Returns the required stake in terms SRI for a given `Balance`.
-async fn required_stake(serai: &TemporalSerai<'_>, balance: ExternalBalance) -> u64 {
-  // This is inclusive to an increase in accuracy
-  let sri_per_coin = serai.dex().oracle_value(balance.coin).await.unwrap().unwrap_or(Amount(0));
-
-  // See dex-pallet for the reasoning on these
-  let coin_decimals = balance.coin.decimals().max(5);
-  let accuracy_increase = u128::from(10u64.pow(coin_decimals));
-
-  let total_coin_value =
-    u64::try_from(u128::from(balance.amount.0) * u128::from(sri_per_coin.0) / accuracy_increase)
-      .unwrap_or(u64::MAX);
-
-  // required stake formula (COIN_VALUE * 1.5) + margin(20%)
-  let required_stake = total_coin_value.saturating_mul(3).saturating_div(2);
-  required_stake.saturating_add(total_coin_value.saturating_div(5))
-}
-
-async fn wait_for_session_change(serai: &Serai) -> u32 {
-  let current_session = serai
-    .as_of_latest_finalized_block()
-    .await
-    .unwrap()
-    .validator_sets()
-    .session(NetworkId::Serai)
-    .await
-    .unwrap()
-    .unwrap()
-    .0;
-  let next_session = current_session + 1;
-
-  // lets wait double the epoch time.
-  tokio::time::timeout(
-    tokio::time::Duration::from_secs(FAST_EPOCH_DURATION * TARGET_BLOCK_TIME * 2),
-    async {
-      while serai
-        .as_of_latest_finalized_block()
-        .await
-        .unwrap()
-        .validator_sets()
-        .session(NetworkId::Serai)
-        .await
-        .unwrap()
-        .unwrap()
-        .0 <
-        next_session
-      {
-        tokio::time::sleep(Duration::from_secs(6)).await;
-      }
-    },
-  )
-  .await
-  .unwrap();
-
-  next_session
-}
-
-async fn get_distances(
-  serai: &TemporalSerai<'_>,
-  current_stake: &HashMap<NetworkId, u64>,
-) -> (HashMap<NetworkId, u64>, u64) {
-  // we should be in the initial period, so calculate how much each network supposedly get..
-  // we can check the supply to see how much coin hence liability we have.
-  let mut distances: HashMap<NetworkId, u64> = HashMap::new();
-  let mut total_distance = 0;
-  for n in EXTERNAL_NETWORKS {
-    let mut required = 0;
-    for c in n.coins() {
-      let amount = serai.coins().coin_supply(c.into()).await.unwrap();
-      required += required_stake(serai, ExternalBalance { coin: c, amount }).await;
-    }
-
-    let mut current = *current_stake.get(&n.into()).unwrap();
-    if current > required {
-      current = required;
-    }
-
-    let distance = required - current;
-    total_distance += distance;
-
-    distances.insert(n.into(), distance);
-  }
-
-  // add serai network portion(20%)
-  let new_total_distance = total_distance.saturating_mul(10) / 8;
-  distances.insert(NetworkId::Serai, new_total_distance - total_distance);
-  total_distance = new_total_distance;
-
-  (distances, total_distance)
-}
-
-async fn get_session_blocks(serai: &TemporalSerai<'_>, session: u32) -> u64 {
-  let begin_block = serai
-    .validator_sets()
-    .session_begin_block(NetworkId::Serai, Session(session))
-    .await
-    .unwrap()
-    .unwrap();
-
-  let next_begin_block = serai
-    .validator_sets()
-    .session_begin_block(NetworkId::Serai, Session(session + 1))
-    .await
-    .unwrap()
-    .unwrap();
-
-  next_begin_block.saturating_sub(begin_block)
-}

substrate/emissions/pallet/Cargo.toml

@@ -39,6 +39,18 @@ serai-primitives = { path = "../../primitives", default-features = false }
 validator-sets-primitives = { package = "serai-validator-sets-primitives", path = "../../validator-sets/primitives", default-features = false }
 emissions-primitives = { package = "serai-emissions-primitives", path = "../primitives", default-features = false }
 
+[dev-dependencies]
+pallet-babe = { git = "https://github.com/serai-dex/substrate", default-features = false }
+pallet-grandpa = { git = "https://github.com/serai-dex/substrate", default-features = false }
+pallet-timestamp = { git = "https://github.com/serai-dex/substrate", default-features = false }
+
+sp-core = { git = "https://github.com/serai-dex/substrate", default-features = false }
+sp-io = { git = "https://github.com/serai-dex/substrate", default-features = false }
+
+serai-abi = { path = "../../abi", default-features = false, features = ["serde"] }
+
+rand_core = "0.6"
+
 [features]
 std = [
   "scale/std",
@@ -49,6 +61,11 @@ std = [
 
   "sp-std/std",
   "sp-runtime/std",
+  "sp-core/std",
+  "sp-io/std",
+
+  "serai-abi/std",
+  "serai-abi/serde",
 
   "coins-pallet/std",
   "validator-sets-pallet/std",
@@ -59,7 +76,19 @@ std = [
 
   "serai-primitives/std",
   "emissions-primitives/std",
+
+  "pallet-babe/std",
+  "pallet-grandpa/std",
+  "pallet-timestamp/std",
 ]
+
+try-runtime = [
+  "frame-system/try-runtime",
+  "frame-support/try-runtime",
+
+  "sp-runtime/try-runtime",
+]
+
 fast-epoch = []
-try-runtime = [] # TODO
+
 default = ["std"]

substrate/emissions/pallet/src/lib.rs

@@ -1,5 +1,11 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
+#[cfg(test)]
+mod tests;
+
+#[cfg(test)]
+mod mock;
+
 #[allow(
   unreachable_patterns,
   clippy::cast_possible_truncation,
@@ -308,15 +314,9 @@ pub mod pallet {
     }
 
     fn initial_period(n: BlockNumberFor<T>) -> bool {
-      #[cfg(feature = "fast-epoch")]
-      let initial_period_duration = FAST_EPOCH_INITIAL_PERIOD;
-
-      #[cfg(not(feature = "fast-epoch"))]
-      let initial_period_duration = 2 * MONTHS;
-
       let genesis_complete_block = GenesisLiquidity::<T>::genesis_complete_block();
       genesis_complete_block.is_some() &&
-        (n.saturated_into::<u64>() < (genesis_complete_block.unwrap() + initial_period_duration))
+        (n.saturated_into::<u64>() < (genesis_complete_block.unwrap() + (2 * MONTHS)))
     }
 
     /// Returns true if any of the external networks haven't reached economic security yet.
@@ -344,17 +344,21 @@ pub mod pallet {
       }
 
       // stake the rewards
-      for (p, score) in scores {
+      let mut total_reward_distributed = 0u64;
-        let p_reward = u64::try_from(
+      for (i, (p, score)) in scores.iter().enumerate() {
-          u128::from(reward).saturating_mul(u128::from(score)) / u128::from(total_score),
+        let p_reward = if i == (scores.len() - 1) {
-        )
+          reward.saturating_sub(total_reward_distributed)
-        .unwrap();
+        } else {
+          u64::try_from(
+            u128::from(reward).saturating_mul(u128::from(*score)) / u128::from(total_score),
+          )
+          .unwrap()
+        };
 
-        Coins::<T>::mint(p, Balance { coin: Coin::Serai, amount: Amount(p_reward) }).unwrap();
+        Coins::<T>::mint(*p, Balance { coin: Coin::Serai, amount: Amount(p_reward) }).unwrap();
-        if ValidatorSets::<T>::distribute_block_rewards(n, p, Amount(p_reward)).is_err() {
+        ValidatorSets::<T>::distribute_block_rewards(n, *p, Amount(p_reward)).unwrap();
-          // TODO: log the failure
+
-          continue;
+        total_reward_distributed = total_reward_distributed.saturating_add(p_reward);
-        }
       }
     }
 

substrate/emissions/pallet/src/mock.rs

@@ -0,0 +1,200 @@
+//! Test environment for Dex pallet.
+
+use super::*;
+
+use std::collections::HashMap;
+
+use frame_support::{
+  construct_runtime,
+  traits::{ConstU16, ConstU32, ConstU64},
+};
+
+use sp_core::{H256, Pair, sr25519::Public};
+use sp_runtime::{
+  traits::{BlakeTwo256, IdentityLookup},
+  BuildStorage,
+};
+
+use serai_primitives::*;
+use validator_sets::{primitives::MAX_KEY_SHARES_PER_SET, MembershipProof};
+
+use crate as emissions;
+pub use coins_pallet as coins;
+pub use validator_sets_pallet as validator_sets;
+pub use genesis_liquidity_pallet as genesis_liquidity;
+pub use dex_pallet as dex;
+pub use pallet_babe as babe;
+pub use pallet_grandpa as grandpa;
+pub use pallet_timestamp as timestamp;
+pub use economic_security_pallet as economic_security;
+
+type Block = frame_system::mocking::MockBlock<Test>;
+// Maximum number of authorities per session.
+pub type MaxAuthorities = ConstU32<{ MAX_KEY_SHARES_PER_SET }>;
+
+pub const MEDIAN_PRICE_WINDOW_LENGTH: u16 = 10;
+
+construct_runtime!(
+  pub enum Test
+  {
+    System: frame_system,
+    Timestamp: timestamp,
+    Coins: coins,
+    LiquidityTokens: coins::<Instance1>::{Pallet, Call, Storage, Event<T>},
+    Emissions: emissions,
+    ValidatorSets: validator_sets,
+    GenesisLiquidity: genesis_liquidity,
+    Dex: dex,
+    Babe: babe,
+    Grandpa: grandpa,
+    EconomicSecurity: economic_security,
+  }
+);
+
+impl frame_system::Config for Test {
+  type BaseCallFilter = frame_support::traits::Everything;
+  type BlockWeights = ();
+  type BlockLength = ();
+  type RuntimeOrigin = RuntimeOrigin;
+  type RuntimeCall = RuntimeCall;
+  type Nonce = u64;
+  type Hash = H256;
+  type Hashing = BlakeTwo256;
+  type AccountId = Public;
+  type Lookup = IdentityLookup<Self::AccountId>;
+  type Block = Block;
+  type RuntimeEvent = RuntimeEvent;
+  type BlockHashCount = ConstU64<250>;
+  type DbWeight = ();
+  type Version = ();
+  type PalletInfo = PalletInfo;
+  type AccountData = ();
+  type OnNewAccount = ();
+  type OnKilledAccount = ();
+  type SystemWeightInfo = ();
+  type SS58Prefix = ();
+  type OnSetCode = ();
+  type MaxConsumers = ConstU32<16>;
+}
+
+impl timestamp::Config for Test {
+  type Moment = u64;
+  type OnTimestampSet = Babe;
+  type MinimumPeriod = ConstU64<{ (TARGET_BLOCK_TIME * 1000) / 2 }>;
+  type WeightInfo = ();
+}
+
+impl babe::Config for Test {
+  type EpochDuration = ConstU64<{ FAST_EPOCH_DURATION }>;
+
+  type ExpectedBlockTime = ConstU64<{ TARGET_BLOCK_TIME * 1000 }>;
+  type EpochChangeTrigger = babe::ExternalTrigger;
+  type DisabledValidators = ValidatorSets;
+
+  type WeightInfo = ();
+  type MaxAuthorities = MaxAuthorities;
+
+  type KeyOwnerProof = MembershipProof<Self>;
+  type EquivocationReportSystem = ();
+}
+
+impl grandpa::Config for Test {
+  type RuntimeEvent = RuntimeEvent;
+
+  type WeightInfo = ();
+  type MaxAuthorities = MaxAuthorities;
+
+  type MaxSetIdSessionEntries = ConstU64<0>;
+  type KeyOwnerProof = MembershipProof<Self>;
+  type EquivocationReportSystem = ();
+}
+
+impl coins::Config for Test {
+  type RuntimeEvent = RuntimeEvent;
+  type AllowMint = ValidatorSets;
+}
+
+impl coins::Config<coins::Instance1> for Test {
+  type RuntimeEvent = RuntimeEvent;
+  type AllowMint = ();
+}
+
+impl dex::Config for Test {
+  type RuntimeEvent = RuntimeEvent;
+
+  type LPFee = ConstU32<3>; // 0.3%
+  type MintMinLiquidity = ConstU64<10000>;
+
+  type MaxSwapPathLength = ConstU32<3>; // coin1 -> SRI -> coin2
+
+  type MedianPriceWindowLength = ConstU16<{ MEDIAN_PRICE_WINDOW_LENGTH }>;
+
+  type WeightInfo = dex::weights::SubstrateWeight<Test>;
+}
+
+impl validator_sets::Config for Test {
+  type RuntimeEvent = RuntimeEvent;
+  type ShouldEndSession = Babe;
+}
+
+impl genesis_liquidity::Config for Test {
+  type RuntimeEvent = RuntimeEvent;
+}
+
+impl economic_security::Config for Test {
+  type RuntimeEvent = RuntimeEvent;
+}
+
+impl Config for Test {
+  type RuntimeEvent = RuntimeEvent;
+}
+
+// Amounts for single key share per network
+pub fn key_shares() -> HashMap<NetworkId, Amount> {
+  HashMap::from([
+    (NetworkId::Serai, Amount(50_000 * 10_u64.pow(8))),
+    (NetworkId::External(ExternalNetworkId::Bitcoin), Amount(1_000_000 * 10_u64.pow(8))),
+    (NetworkId::External(ExternalNetworkId::Ethereum), Amount(1_000_000 * 10_u64.pow(8))),
+    (NetworkId::External(ExternalNetworkId::Monero), Amount(100_000 * 10_u64.pow(8))),
+  ])
+}
+
+pub(crate) fn new_test_ext() -> sp_io::TestExternalities {
+  let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
+  let networks: Vec<(NetworkId, Amount)> = key_shares().into_iter().collect::<Vec<_>>();
+
+  let accounts: Vec<Public> = vec![
+    insecure_pair_from_name("Alice").public(),
+    insecure_pair_from_name("Bob").public(),
+    insecure_pair_from_name("Charlie").public(),
+    insecure_pair_from_name("Dave").public(),
+    insecure_pair_from_name("Eve").public(),
+    insecure_pair_from_name("Ferdie").public(),
+  ];
+  let validators = accounts.clone();
+
+  coins::GenesisConfig::<Test> {
+    accounts: accounts
+      .into_iter()
+      .map(|a| (a, Balance { coin: Coin::Serai, amount: Amount(1 << 60) }))
+      .collect(),
+    _ignore: Default::default(),
+  }
+  .assimilate_storage(&mut t)
+  .unwrap();
+
+  validator_sets::GenesisConfig::<Test> {
+    networks: networks.clone(),
+    participants: validators.clone(),
+  }
+  .assimilate_storage(&mut t)
+  .unwrap();
+
+  crate::GenesisConfig::<Test> { networks, participants: validators.clone() }
+    .assimilate_storage(&mut t)
+    .unwrap();
+
+  let mut ext = sp_io::TestExternalities::new(t);
+  ext.execute_with(|| System::set_block_number(0));
+  ext
+}

substrate/emissions/pallet/src/tests.rs

@@ -0,0 +1,467 @@
+use crate::{mock::*, primitives::*};
+
+use std::collections::HashMap;
+
+use rand_core::{RngCore, OsRng};
+
+use sp_core::{sr25519::Signature, Pair};
+use sp_runtime::BoundedVec;
+
+use frame_system::RawOrigin;
+use frame_support::traits::{Hooks, Get};
+
+use genesis_liquidity_pallet::{
+  Pallet as GenesisLiquidity,
+  primitives::{Values, GENESIS_LIQUIDITY_ACCOUNT},
+};
+use validator_sets_pallet::{Pallet as ValidatorSets, primitives::Session};
+use coins_pallet::Pallet as Coins;
+use dex_pallet::Pallet as Dex;
+use economic_security::Pallet as EconomicSecurity;
+
+use serai_primitives::*;
+use validator_sets_primitives::{KeyPair, ValidatorSet};
+
+fn set_up_genesis() -> u64 {
+  // make accounts with amounts
+  let mut accounts = HashMap::new();
+  for coin in EXTERNAL_COINS {
+    // make 5 accounts per coin
+    let mut values = vec![];
+    for _ in 0 .. 5 {
+      let mut address = SeraiAddress::new([0; 32]);
+      OsRng.fill_bytes(&mut address.0);
+      values.push((address, Amount(OsRng.next_u64() % (10_000 * 10u64.pow(coin.decimals())))));
+    }
+    accounts.insert(coin, values);
+  }
+
+  // add some genesis liquidity
+  for (coin, amounts) in accounts {
+    for (address, amount) in amounts {
+      let balance = ExternalBalance { coin, amount };
+
+      Coins::<Test>::mint(GENESIS_LIQUIDITY_ACCOUNT.into(), balance.into()).unwrap();
+      GenesisLiquidity::<Test>::add_coin_liquidity(address.into(), balance).unwrap();
+    }
+  }
+
+  // make genesis liquidity event happen
+  let block_number = MONTHS;
+  let values = Values { monero: 184100, ether: 4785000, dai: 1500 };
+  GenesisLiquidity::<Test>::oraclize_values(RawOrigin::None.into(), values, Signature([0u8; 64]))
+    .unwrap();
+  GenesisLiquidity::<Test>::on_initialize(block_number);
+  System::set_block_number(block_number);
+
+  // populate the coin values
+  Dex::<Test>::on_finalize(block_number);
+
+  block_number
+}
+
+// TODO: make this fn belong to the pallet itself use it there as well?
+// The problem with that would be if there is a problem with this function
+// tests can't catch it since it would the same fn?
+fn distances() -> (HashMap<NetworkId, u64>, u64) {
+  let mut distances = HashMap::new();
+  let mut total_distance: u64 = 0;
+
+  // calculate distance to economic security per network
+  for n in EXTERNAL_NETWORKS {
+    let required = ValidatorSets::<Test>::required_stake_for_network(n);
+    let mut current =
+      ValidatorSets::<Test>::total_allocated_stake(NetworkId::from(n)).unwrap_or(Amount(0)).0;
+    if current > required {
+      current = required;
+    }
+
+    let distance = required - current;
+    distances.insert(n.into(), distance);
+    total_distance = total_distance.saturating_add(distance);
+  }
+
+  // add serai network portion (20%)
+  let new_total_distance = total_distance.saturating_mul(100) / (100 - 20);
+  distances.insert(NetworkId::Serai, new_total_distance - total_distance);
+  total_distance = new_total_distance;
+
+  (distances, total_distance)
+}
+
+fn set_keys_for_session() {
+  for network in EXTERNAL_NETWORKS {
+    ValidatorSets::<Test>::set_keys(
+      RawOrigin::None.into(),
+      network,
+      BoundedVec::new(),
+      KeyPair(insecure_pair_from_name("Alice").public(), vec![].try_into().unwrap()),
+      Signature([0u8; 64]),
+    )
+    .unwrap();
+  }
+}
+
+fn make_fake_swap_volume() {
+  let acc = insecure_pair_from_name("random").public();
+  for _ in 0 .. 10 {
+    let path_len = (OsRng.next_u32() % 2) + 2;
+
+    let coins = &COINS[1 ..];
+    let path = if path_len == 2 {
+      let coin = coins[(OsRng.next_u32() as usize) % coins.len()];
+      let in_or_out = (OsRng.next_u32() % 2) == 0;
+      if in_or_out {
+        vec![coin, Coin::Serai]
+      } else {
+        vec![Coin::Serai, coin]
+      }
+    } else {
+      let in_coin = coins[(OsRng.next_u32() as usize) % coins.len()];
+      let coins_without_in_coin = coins.iter().filter(|&c| *c != in_coin).collect::<Vec<_>>();
+      let out_coin =
+        coins_without_in_coin[(OsRng.next_u32() as usize) % coins_without_in_coin.len()];
+      vec![in_coin, Coin::Serai, *out_coin]
+    };
+
+    let one_in_coin = 10u64.pow(path[0].decimals());
+    Coins::<Test>::mint(acc, Balance { coin: path[0], amount: Amount(2 * one_in_coin) }).unwrap();
+    let amount_in = OsRng.next_u64() % (one_in_coin);
+
+    Dex::<Test>::swap_exact_tokens_for_tokens(
+      RawOrigin::Signed(acc).into(),
+      path.try_into().unwrap(),
+      amount_in,
+      1,
+      acc,
+    )
+    .unwrap();
+  }
+}
+
+fn get_session_swap_volumes(
+  last_swap_volume: &mut HashMap<ExternalCoin, u64>,
+) -> (HashMap<ExternalCoin, u64>, HashMap<NetworkId, u64>, u64) {
+  let mut volume_per_coin: HashMap<ExternalCoin, u64> = HashMap::new();
+  for c in EXTERNAL_COINS {
+    let current_volume = Dex::<Test>::swap_volume(c).unwrap_or(0);
+    let last_volume = last_swap_volume.get(&c).unwrap_or(&0);
+    let vol_this_epoch = current_volume.saturating_sub(*last_volume);
+
+    // update the current volume
+    last_swap_volume.insert(c, current_volume);
+    volume_per_coin.insert(c, vol_this_epoch);
+  }
+
+  // aggregate per network
+  let mut total_volume = 0u64;
+  let mut volume_per_network: HashMap<NetworkId, u64> = HashMap::new();
+  for (c, vol) in &volume_per_coin {
+    volume_per_network.insert(
+      c.network().into(),
+      (*volume_per_network.get(&c.network().into()).unwrap_or(&0)).saturating_add(*vol),
+    );
+    total_volume = total_volume.saturating_add(*vol);
+  }
+  volume_per_network.insert(NetworkId::Serai, 0);
+
+  (volume_per_coin, volume_per_network, total_volume)
+}
+
+fn get_pool_vs_validator_rewards(n: NetworkId, reward: u64) -> (u64, u64) {
+  if n == NetworkId::Serai {
+    (reward, 0)
+  } else {
+    // calculate pool vs validator share
+    let capacity = ValidatorSets::<Test>::total_allocated_stake(n).unwrap_or(Amount(0)).0;
+    let required = ValidatorSets::<Test>::required_stake_for_network(n.try_into().unwrap());
+    let unused_capacity = capacity.saturating_sub(required);
+
+    let distribution = unused_capacity.saturating_mul(ACCURACY_MULTIPLIER) / capacity;
+    let total = DESIRED_DISTRIBUTION.saturating_add(distribution);
+
+    let validators_reward = DESIRED_DISTRIBUTION.saturating_mul(reward) / total;
+    let network_pool_reward = reward.saturating_sub(validators_reward);
+    (validators_reward, network_pool_reward)
+  }
+}
+
+#[test]
+fn check_pre_ec_security_initial_period_emissions() {
+  new_test_ext().execute_with(|| {
+    // set up genesis liquidity
+    let mut block_number = set_up_genesis();
+
+    for _ in 0 .. 5 {
+      // set session keys. we need this here before reading the current stakes for session 0.
+      // We need it for other sessions to be able to retire the set.
+      set_keys_for_session();
+
+      // get current stakes
+      let mut current_stake = HashMap::new();
+      for n in NETWORKS {
+        current_stake.insert(n, ValidatorSets::<Test>::total_allocated_stake(n).unwrap().0);
+      }
+
+      // trigger rewards distribution for the past session
+      ValidatorSets::<Test>::new_session();
+      Emissions::on_initialize(block_number + 1);
+
+      // calculate the total reward for this epoch
+      let (distances, total_distance) = distances();
+      let session = ValidatorSets::<Test>::session(NetworkId::Serai).unwrap_or(Session(0));
+      let block_count = ValidatorSets::<Test>::session_begin_block(NetworkId::Serai, session) -
+        ValidatorSets::<Test>::session_begin_block(NetworkId::Serai, Session(session.0 - 1));
+      let reward_this_epoch = block_count * INITIAL_REWARD_PER_BLOCK;
+
+      let reward_per_network = distances
+        .into_iter()
+        .map(|(n, distance)| {
+          // calculate how much each network gets based on distance to ec-security
+          let reward = u64::try_from(
+            u128::from(reward_this_epoch).saturating_mul(u128::from(distance)) /
+              u128::from(total_distance),
+          )
+          .unwrap();
+          (n, reward)
+        })
+        .collect::<HashMap<NetworkId, u64>>();
+
+      for (n, reward) in reward_per_network {
+        ValidatorSets::<Test>::retire_set(ValidatorSet {
+          session: Session(session.0 - 1),
+          network: n,
+        });
+
+        // all validator rewards should automatically be staked
+        assert_eq!(
+          ValidatorSets::<Test>::total_allocated_stake(n).unwrap().0,
+          *current_stake.get(&n).unwrap() + reward
+        );
+      }
+
+      block_number += <<Test as pallet_babe::Config>::EpochDuration as Get<u64>>::get();
+      System::set_block_number(block_number);
+    }
+  });
+}
+
+#[test]
+fn check_pre_ec_security_emissions() {
+  new_test_ext().execute_with(|| {
+    // set up genesis liquidity
+    let mut block_number = set_up_genesis();
+
+    // move the block number out of initial period which is 2 more months
+    block_number += 2 * MONTHS;
+    System::set_block_number(block_number);
+
+    // make a fresh session
+    set_keys_for_session();
+    ValidatorSets::<Test>::new_session();
+    for network in NETWORKS {
+      ValidatorSets::<Test>::retire_set(ValidatorSet { session: Session(0), network });
+    }
+
+    // move the block for the next session
+    block_number += <<Test as pallet_babe::Config>::EpochDuration as Get<u64>>::get();
+    System::set_block_number(block_number);
+
+    for _ in 0 .. 5 {
+      // set session keys. we need this here before reading the current stakes for session 0.
+      // We need it for other sessions to be able to retire the set.
+      set_keys_for_session();
+
+      // get current stakes
+      let mut current_stake = HashMap::new();
+      for n in NETWORKS {
+        current_stake.insert(n, ValidatorSets::<Test>::total_allocated_stake(n).unwrap().0);
+      }
+
+      // trigger rewards distribution for the past session
+      ValidatorSets::<Test>::new_session();
+      <Emissions as Hooks<BlockNumber>>::on_initialize(block_number + 1);
+
+      // calculate the total reward for this epoch
+      let (distances, total_distance) = distances();
+      let session = ValidatorSets::<Test>::session(NetworkId::Serai).unwrap_or(Session(0));
+      let block_count = ValidatorSets::<Test>::session_begin_block(NetworkId::Serai, session) -
+        ValidatorSets::<Test>::session_begin_block(NetworkId::Serai, Session(session.0 - 1));
+      let reward_this_epoch = block_count * (total_distance / (SECURE_BY - block_number));
+
+      let reward_per_network = distances
+        .into_iter()
+        .map(|(n, distance)| {
+          // calculate how much each network gets based on distance to ec-security
+          let reward = u64::try_from(
+            u128::from(reward_this_epoch).saturating_mul(u128::from(distance)) /
+              u128::from(total_distance),
+          )
+          .unwrap();
+          (n, reward)
+        })
+        .collect::<HashMap<NetworkId, u64>>();
+
+      for (n, reward) in reward_per_network {
+        ValidatorSets::<Test>::retire_set(ValidatorSet {
+          session: Session(session.0 - 1),
+          network: n,
+        });
+
+        // all validator rewards should automatically be staked
+        assert_eq!(
+          ValidatorSets::<Test>::total_allocated_stake(n).unwrap().0,
+          *current_stake.get(&n).unwrap() + reward
+        );
+      }
+
+      block_number += <<Test as pallet_babe::Config>::EpochDuration as Get<u64>>::get();
+      System::set_block_number(block_number);
+    }
+  });
+}
+
+#[test]
+fn check_post_ec_security_emissions() {
+  new_test_ext().execute_with(|| {
+    // set up genesis liquidity
+    let mut block_number = set_up_genesis();
+
+    // make all networks reach economic security
+    set_keys_for_session();
+    let (distances, _) = distances();
+    for (network, distance) in distances {
+      if network == NetworkId::Serai {
+        continue;
+      }
+
+      let participants =
+        ValidatorSets::<Test>::participants_for_latest_decided_set(network).unwrap();
+      let al_per_key_share = ValidatorSets::<Test>::allocation_per_key_share(network).unwrap().0;
+
+      // we want some unused capacity so we stake more SRI than necessary
+      let mut key_shares = (distance / al_per_key_share) + 10;
+
+      'outer: while key_shares > 0 {
+        for (account, _) in &participants {
+          ValidatorSets::<Test>::distribute_block_rewards(
+            network,
+            *account,
+            Amount(al_per_key_share),
+          )
+          .unwrap();
+
+          if key_shares > 0 {
+            key_shares -= 1;
+          } else {
+            break 'outer;
+          }
+        }
+      }
+    }
+
+    // update TAS
+    ValidatorSets::<Test>::new_session();
+    for network in NETWORKS {
+      ValidatorSets::<Test>::retire_set(ValidatorSet { session: Session(0), network });
+    }
+
+    // make sure we reached economic security
+    <EconomicSecurity<Test> as Hooks<BlockNumber>>::on_initialize(block_number);
+    for n in EXTERNAL_NETWORKS {
+      EconomicSecurity::<Test>::economic_security_block(n).unwrap();
+    }
+
+    // move the block number for the next session
+    block_number += <<Test as pallet_babe::Config>::EpochDuration as Get<u64>>::get();
+    System::set_block_number(block_number);
+
+    let mut last_swap_volume = HashMap::new();
+    for _ in 0 .. 5 {
+      set_keys_for_session();
+
+      // make some fake swap volume
+      make_fake_swap_volume();
+      let (vpc, vpn, total_volume) = get_session_swap_volumes(&mut last_swap_volume);
+
+      // get current stakes & each pool SRI amounts
+      let mut current_stake = HashMap::new();
+      let mut current_pool_coins = HashMap::new();
+      for n in NETWORKS {
+        current_stake.insert(n, ValidatorSets::<Test>::total_allocated_stake(n).unwrap().0);
+        if let NetworkId::External(network) = n {
+          for c in network.coins() {
+            let acc = Dex::<Test>::get_pool_account(c);
+            current_pool_coins.insert(c, Coins::<Test>::balance(acc, Coin::Serai).0);
+          }
+        }
+      }
+
+      // trigger rewards distribution for the past session
+      ValidatorSets::<Test>::new_session();
+      <Emissions as Hooks<BlockNumber>>::on_initialize(block_number + 1);
+
+      // calculate the total reward for this epoch
+      let session = ValidatorSets::<Test>::session(NetworkId::Serai).unwrap_or(Session(0));
+      let block_count = ValidatorSets::<Test>::session_begin_block(NetworkId::Serai, session) -
+        ValidatorSets::<Test>::session_begin_block(NetworkId::Serai, Session(session.0 - 1));
+      let reward_this_epoch = block_count * REWARD_PER_BLOCK;
+
+      let reward_per_network = vpn
+        .iter()
+        .map(|(n, volume)| {
+          let reward = if *n == NetworkId::Serai {
+            reward_this_epoch / 5
+          } else {
+            let reward = reward_this_epoch - (reward_this_epoch / 5);
+            // TODO: It is highly unlikely but what to do in case of 0 total volume?
+            if total_volume != 0 {
+              u64::try_from(
+                u128::from(reward).saturating_mul(u128::from(*volume)) / u128::from(total_volume),
+              )
+              .unwrap()
+            } else {
+              0
+            }
+          };
+          (*n, reward)
+        })
+        .collect::<HashMap<NetworkId, u64>>();
+
+      for (n, reward) in reward_per_network {
+        let (validator_rewards, network_pool_rewards) = get_pool_vs_validator_rewards(n, reward);
+        ValidatorSets::<Test>::retire_set(ValidatorSet {
+          session: Session(session.0 - 1),
+          network: n,
+        });
+
+        // all validator rewards should automatically be staked
+        assert_eq!(
+          ValidatorSets::<Test>::total_allocated_stake(n).unwrap().0,
+          *current_stake.get(&n).unwrap() + validator_rewards
+        );
+
+        // all pool rewards should be available in the pool account
+        if network_pool_rewards != 0 {
+          for coin in n.coins() {
+            let c: ExternalCoin = coin.try_into().unwrap();
+            let pool_reward = u64::try_from(
+              u128::from(network_pool_rewards).saturating_mul(u128::from(vpc[&c])) /
+                u128::from(vpn[&n]),
+            )
+            .unwrap();
+
+            let acc = Dex::<Test>::get_pool_account(c);
+            assert_eq!(
+              Coins::<Test>::balance(acc, Coin::Serai).0,
+              current_pool_coins[&c] + pool_reward
+            )
+          }
+        }
+      }
+
+      block_number += <<Test as pallet_babe::Config>::EpochDuration as Get<u64>>::get();
+      System::set_block_number(block_number);
+    }
+  });
+}

substrate/validator-sets/pallet/src/lib.rs

@@ -393,6 +393,7 @@ pub mod pallet {
       let allocation_per_key_share = Self::allocation_per_key_share(network).unwrap().0;
 
       let mut participants = vec![];
+      let mut total_allocated_stake = 0;
       {
         let mut iter = SortedAllocationsIter::<T>::new(network);
         let mut key_shares = 0;
@@ -403,6 +404,7 @@ pub mod pallet {
             (amount.0 / allocation_per_key_share).min(u64::from(MAX_KEY_SHARES_PER_SET));
           participants.push((key, these_key_shares));
 
+          total_allocated_stake += amount.0;
           key_shares += these_key_shares;
         }
         amortize_excess_key_shares(&mut participants);
@@ -416,6 +418,10 @@ pub mod pallet {
       Pallet::<T>::deposit_event(Event::NewSet { set });
 
       Participants::<T>::set(network, Some(participants.try_into().unwrap()));
+      if network == NetworkId::Serai {
+        TotalAllocatedStake::<T>::set(network, Some(Amount(total_allocated_stake)));
+      }
+
       SessionBeginBlock::<T>::set(
         network,
         session,
@@ -713,7 +719,7 @@ pub mod pallet {
         }))
     }
 
-    fn new_session() {
+    pub fn new_session() {
       for network in serai_primitives::NETWORKS {
         // If this network hasn't started sessions yet, don't start one now
         let Some(current_session) = Self::session(network) else { continue };