Implement block emissions (#551)

* add genesis liquidity implementation

* add missing deposit event

* fix CI issues

* minor fixes

* make math safer

* fix fmt

* implement block emissions

* make remove liquidity an authorized call

* implement setting initial values for coins

* add genesis liquidity test & misc fixes

* updato develop latest

* fix rotation test

* fix licencing

* add fast-epoch feature

* only create the pool when adding liquidity first time

* add initial reward era test

* test whole pre ec security emissions

* fix clippy

* add swap-to-staked-sri feature

* rebase changes

* fix tests

* Remove accidentally commited ETH ABI files

* fix some pr comments

* Finish up fixing pr comments

* exclude SRI from is_allowed check

* Misc changes

---------

Co-authored-by: akildemir <aeg_asd@hotmail.com>
Co-authored-by: Luke Parker <lukeparker5132@gmail.com>

substrate/client/tests/common/genesis_liquidity.rs

@@ -0,0 +1,115 @@
+use std::collections::HashMap;
+
+use rand_core::{RngCore, OsRng};
+use zeroize::Zeroizing;
+
+use ciphersuite::{Ciphersuite, Ristretto};
+use frost::dkg::musig::musig;
+use schnorrkel::Schnorrkel;
+
+use sp_core::{sr25519::Signature, Pair as PairTrait};
+
+use serai_abi::{
+  genesis_liquidity::primitives::{oraclize_values_message, Values},
+  validator_sets::primitives::{musig_context, Session, ValidatorSet},
+  in_instructions::primitives::{InInstruction, InInstructionWithBalance, Batch},
+  primitives::{
+    Amount, NetworkId, Coin, Balance, BlockHash, SeraiAddress, insecure_pair_from_name,
+  },
+};
+
+use serai_client::{Serai, SeraiGenesisLiquidity};
+
+use crate::common::{in_instructions::provide_batch, tx::publish_tx};
+
+#[allow(dead_code)]
+pub async fn set_up_genesis(
+  serai: &Serai,
+  coins: &[Coin],
+  values: &HashMap<Coin, u64>,
+) -> (HashMap<Coin, Vec<(SeraiAddress, Amount)>>, HashMap<NetworkId, u32>) {
+  // make accounts with amounts
+  let mut accounts = HashMap::new();
+  for coin in 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() % 10u64.pow(coin.decimals()))));
+    }
+    accounts.insert(*coin, values);
+  }
+
+  // send a batch per coin
+  let mut batch_ids: HashMap<NetworkId, u32> = HashMap::new();
+  for coin in coins {
+    // set up instructions
+    let instructions = accounts[coin]
+      .iter()
+      .map(|(addr, amount)| InInstructionWithBalance {
+        instruction: InInstruction::GenesisLiquidity(*addr),
+        balance: Balance { coin: *coin, amount: *amount },
+      })
+      .collect::<Vec<_>>();
+
+    // set up bloch hash
+    let mut block = BlockHash([0; 32]);
+    OsRng.fill_bytes(&mut block.0);
+
+    // set up batch id
+    batch_ids
+      .entry(coin.network())
+      .and_modify(|v| {
+        *v += 1;
+      })
+      .or_insert(0);
+
+    let batch =
+      Batch { network: coin.network(), id: batch_ids[&coin.network()], block, instructions };
+    provide_batch(serai, batch).await;
+  }
+
+  // set values relative to each other. We can do that without checking for genesis period blocks
+  // since we are running in test(fast-epoch) mode.
+  // TODO: Random values here
+  let values =
+    Values { monero: values[&Coin::Monero], ether: values[&Coin::Ether], dai: values[&Coin::Dai] };
+  set_values(serai, &values).await;
+
+  (accounts, batch_ids)
+}
+
+#[allow(dead_code)]
+async fn set_values(serai: &Serai, values: &Values) {
+  // prepare a Musig tx to oraclize the relative values
+  let pair = insecure_pair_from_name("Alice");
+  let public = pair.public();
+  // we publish the tx in set 1
+  let set = ValidatorSet { session: Session(1), network: NetworkId::Serai };
+
+  let public_key = <Ristretto as Ciphersuite>::read_G::<&[u8]>(&mut public.0.as_ref()).unwrap();
+  let secret_key = <Ristretto as Ciphersuite>::read_F::<&[u8]>(
+    &mut pair.as_ref().secret.to_bytes()[.. 32].as_ref(),
+  )
+  .unwrap();
+
+  assert_eq!(Ristretto::generator() * secret_key, public_key);
+  let threshold_keys =
+    musig::<Ristretto>(&musig_context(set), &Zeroizing::new(secret_key), &[public_key]).unwrap();
+
+  let sig = frost::tests::sign_without_caching(
+    &mut OsRng,
+    frost::tests::algorithm_machines(
+      &mut OsRng,
+      &Schnorrkel::new(b"substrate"),
+      &HashMap::from([(threshold_keys.params().i(), threshold_keys.into())]),
+    ),
+    &oraclize_values_message(&set, values),
+  );
+
+  // oraclize values
+  let _ =
+    publish_tx(serai, &SeraiGenesisLiquidity::oraclize_values(*values, Signature(sig.to_bytes())))
+      .await;
+}

substrate/client/tests/common/in_instructions.rs

@@ -10,7 +10,7 @@ use sp_core::Pair;
 
 use serai_client::{
   primitives::{insecure_pair_from_name, BlockHash, NetworkId, Balance, SeraiAddress},
-  validator_sets::primitives::{Session, ValidatorSet, KeyPair},
+  validator_sets::primitives::{ValidatorSet, KeyPair},
   in_instructions::{
     primitives::{Batch, SignedBatch, batch_message, InInstruction, InInstructionWithBalance},
     InInstructionsEvent,
@@ -22,12 +22,12 @@ use crate::common::{tx::publish_tx, validator_sets::set_keys};
 
 #[allow(dead_code)]
 pub async fn provide_batch(serai: &Serai, batch: Batch) -> [u8; 32] {
-  // TODO: Get the latest session
-  let set = ValidatorSet { session: Session(0), network: batch.network };
+  let serai_latest = serai.as_of_latest_finalized_block().await.unwrap();
+  let session = serai_latest.validator_sets().session(batch.network).await.unwrap().unwrap();
+  let set = ValidatorSet { session, network: batch.network };
+
   let pair = insecure_pair_from_name(&format!("ValidatorSet {set:?}"));
-  let keys = if let Some(keys) =
-    serai.as_of_latest_finalized_block().await.unwrap().validator_sets().keys(set).await.unwrap()
-  {
+  let keys = if let Some(keys) = serai_latest.validator_sets().keys(set).await.unwrap() {
     keys
   } else {
     let keys = KeyPair(pair.public(), vec![].try_into().unwrap());

substrate/client/tests/common/mod.rs

@@ -2,6 +2,7 @@ pub mod tx;
 pub mod validator_sets;
 pub mod in_instructions;
 pub mod dex;
+pub mod genesis_liquidity;
 
 #[macro_export]
 macro_rules! serai_test {

substrate/client/tests/dex.rs

@@ -1,14 +1,13 @@
 use rand_core::{RngCore, OsRng};
-use scale::Encode;
 
-use sp_core::{Pair as PairTrait, bounded_vec::BoundedVec, hashing::blake2_256};
+use sp_core::{Pair as PairTrait, bounded_vec::BoundedVec};
 
 use serai_abi::in_instructions::primitives::DexCall;
 
 use serai_client::{
   primitives::{
     Amount, NetworkId, Coin, Balance, BlockHash, insecure_pair_from_name, ExternalAddress,
-    SeraiAddress, PublicKey,
+    SeraiAddress,
   },
   in_instructions::primitives::{
     InInstruction, InInstructionWithBalance, Batch, IN_INSTRUCTION_EXECUTOR, OutAddress,
@@ -28,33 +27,6 @@ use common::{
 // TODO: Modularize common code
 // TODO: Check Transfer events
 serai_test!(
-  create_pool: (|serai: Serai| async move {
-    let block = serai.finalized_block_by_number(0).await.unwrap().unwrap().hash();
-    let events = serai.as_of(block).dex().events().await.unwrap();
-
-    assert_eq!(
-      events,
-      vec![
-        DexEvent::PoolCreated {
-          pool_id: Coin::Bitcoin,
-          pool_account: PublicKey::from_raw(blake2_256(&Coin::Bitcoin.encode())).into(),
-        },
-        DexEvent::PoolCreated {
-          pool_id: Coin::Ether,
-          pool_account: PublicKey::from_raw(blake2_256(&Coin::Ether.encode())).into(),
-        },
-        DexEvent::PoolCreated {
-          pool_id: Coin::Dai,
-          pool_account: PublicKey::from_raw(blake2_256(&Coin::Dai.encode())).into(),
-        },
-        DexEvent::PoolCreated {
-          pool_id: Coin::Monero,
-          pool_account: PublicKey::from_raw(blake2_256(&Coin::Monero.encode())).into(),
-        },
-      ]
-    );
-  })
-
   add_liquidity: (|serai: Serai| async move {
     let coin = Coin::Monero;
     let pair = insecure_pair_from_name("Ferdie");

substrate/client/tests/emissions.rs

@@ -0,0 +1,257 @@
+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, Coin, COINS, FAST_EPOCH_DURATION, FAST_EPOCH_INITIAL_PERIOD, NETWORKS,
+    TARGET_BLOCK_TIME,
+  },
+  validator_sets::primitives::Session,
+};
+
+use serai_client::{
+  primitives::{Amount, NetworkId, Balance},
+  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<NetworkId, u32>) {
+  for network in NETWORKS {
+    if network != NetworkId::Serai {
+      // 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 coins = COINS.into_iter().filter(|c| *c != Coin::native()).collect::<Vec<_>>();
+  let values = HashMap::from([(Coin::Monero, 184100), (Coin::Ether, 4785000), (Coin::Dai, 1500)]);
+  let (_, mut batch_ids) = set_up_genesis(&serai, &coins, &values).await;
+
+  // wait until genesis is complete
+  while !serai
+    .as_of_latest_finalized_block()
+    .await
+    .unwrap()
+    .genesis_liquidity()
+    .genesis_complete()
+    .await
+    .unwrap()
+  {
+    tokio::time::sleep(Duration::from_secs(1)).await;
+  }
+  let genesis_complete_block = serai.latest_finalized_block().await.unwrap().number();
+
+  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 + 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: Balance) -> 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 NETWORKS {
+    if n == NetworkId::Serai {
+      continue;
+    }
+
+    let mut required = 0;
+    for c in n.coins() {
+      let amount = serai.coins().coin_supply(*c).await.unwrap();
+      required += required_stake(serai, Balance { coin: *c, amount }).await;
+    }
+
+    let mut current = *current_stake.get(&n).unwrap();
+    if current > required {
+      current = required;
+    }
+
+    let distance = required - current;
+    total_distance += distance;
+
+    distances.insert(n, 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/client/tests/genesis_liquidity.rs

@@ -1,37 +1,15 @@
 use std::{time::Duration, collections::HashMap};
 
-use rand_core::{RngCore, OsRng};
-use zeroize::Zeroizing;
+use serai_client::Serai;
 
-use ciphersuite::{Ciphersuite, Ristretto};
-use frost::dkg::musig::musig;
-use schnorrkel::Schnorrkel;
+use serai_abi::primitives::{Coin, COINS, Amount, GENESIS_SRI};
 
-use serai_client::{
-  genesis_liquidity::{
-    primitives::{GENESIS_LIQUIDITY_ACCOUNT, INITIAL_GENESIS_LP_SHARES},
-    SeraiGenesisLiquidity,
-  },
-  validator_sets::primitives::{musig_context, Session, ValidatorSet},
-};
-
-use serai_abi::{
-  genesis_liquidity::primitives::{oraclize_values_message, Values},
-  primitives::COINS,
-};
-
-use sp_core::{sr25519::Signature, Pair as PairTrait};
-
-use serai_client::{
-  primitives::{
-    Amount, NetworkId, Coin, Balance, BlockHash, SeraiAddress, insecure_pair_from_name, GENESIS_SRI,
-  },
-  in_instructions::primitives::{InInstruction, InInstructionWithBalance, Batch},
-  Serai,
+use serai_client::genesis_liquidity::primitives::{
+  GENESIS_LIQUIDITY_ACCOUNT, INITIAL_GENESIS_LP_SHARES,
 };
 
 mod common;
-use common::{in_instructions::provide_batch, tx::publish_tx};
+use common::genesis_liquidity::set_up_genesis;
 
 serai_test_fast_epoch!(
   genesis_liquidity: (|serai: Serai| async move {
@@ -39,79 +17,25 @@ serai_test_fast_epoch!(
   })
 );
 
-async fn test_genesis_liquidity(serai: Serai) {
-  // all coins except the native
+pub async fn test_genesis_liquidity(serai: Serai) {
+  // set up the genesis
   let coins = COINS.into_iter().filter(|c| *c != Coin::native()).collect::<Vec<_>>();
+  let values = HashMap::from([(Coin::Monero, 184100), (Coin::Ether, 4785000), (Coin::Dai, 1500)]);
+  let (accounts, _) = set_up_genesis(&serai, &coins, &values).await;
 
-  // make accounts with amounts
-  let mut accounts = HashMap::new();
-  for coin in coins.clone() {
-    // 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() % 10u64.pow(coin.decimals()))));
-    }
-    accounts.insert(coin, values);
+  // wait until genesis is complete
+  while !serai
+    .as_of_latest_finalized_block()
+    .await
+    .unwrap()
+    .genesis_liquidity()
+    .genesis_complete()
+    .await
+    .unwrap()
+  {
+    tokio::time::sleep(Duration::from_secs(1)).await;
   }
 
-  // send a batch per coin
-  let mut batch_ids: HashMap<NetworkId, u32> = HashMap::new();
-  for coin in coins.clone() {
-    // set up instructions
-    let instructions = accounts[&coin]
-      .iter()
-      .map(|(addr, amount)| InInstructionWithBalance {
-        instruction: InInstruction::GenesisLiquidity(*addr),
-        balance: Balance { coin, amount: *amount },
-      })
-      .collect::<Vec<_>>();
-
-    // set up bloch hash
-    let mut block = BlockHash([0; 32]);
-    OsRng.fill_bytes(&mut block.0);
-
-    // set up batch id
-    batch_ids
-      .entry(coin.network())
-      .and_modify(|v| {
-        *v += 1;
-      })
-      .or_insert(0);
-
-    let batch =
-      Batch { network: coin.network(), id: batch_ids[&coin.network()], block, instructions };
-    provide_batch(&serai, batch).await;
-  }
-
-  // wait until genesis ends
-  let genesis_blocks = 10; // TODO
-  let block_time = 6; // TODO
-  tokio::time::timeout(
-    tokio::time::Duration::from_secs(3 * (genesis_blocks * block_time)),
-    async {
-      while serai.latest_finalized_block().await.unwrap().number() < 10 {
-        tokio::time::sleep(Duration::from_secs(6)).await;
-      }
-    },
-  )
-  .await
-  .unwrap();
-
-  // set values relative to each other
-  // TODO: Random values here
-  let values = Values { monero: 184100, ether: 4785000, dai: 1500 };
-  set_values(&serai, &values).await;
-  let values_map = HashMap::from([
-    (Coin::Monero, values.monero),
-    (Coin::Ether, values.ether),
-    (Coin::Dai, values.dai),
-  ]);
-
-  // wait a little bit..
-  tokio::time::sleep(Duration::from_secs(12)).await;
-
   // check total SRI supply is +100M
   // there are 6 endowed accounts in dev-net. Take this into consideration when checking
   // for the total sri minted at this time.
@@ -133,7 +57,7 @@ async fn test_genesis_liquidity(serai: Serai) {
   for coin in coins.clone() {
     let total_coin = accounts[&coin].iter().fold(0u128, |acc, value| acc + u128::from(value.1 .0));
     let value = if coin != Coin::Bitcoin {
-      (total_coin * u128::from(values_map[&coin])) / 10u128.pow(coin.decimals())
+      (total_coin * u128::from(values[&coin])) / 10u128.pow(coin.decimals())
     } else {
       total_coin
     };
@@ -181,36 +105,3 @@ async fn test_genesis_liquidity(serai: Serai) {
 
   // TODO: test remove the liq before/after genesis ended.
 }
-
-async fn set_values(serai: &Serai, values: &Values) {
-  // prepare a Musig tx to oraclize the relative values
-  let pair = insecure_pair_from_name("Alice");
-  let public = pair.public();
-  // we publish the tx in set 4
-  let set = ValidatorSet { session: Session(4), network: NetworkId::Serai };
-
-  let public_key = <Ristretto as Ciphersuite>::read_G::<&[u8]>(&mut public.0.as_ref()).unwrap();
-  let secret_key = <Ristretto as Ciphersuite>::read_F::<&[u8]>(
-    &mut pair.as_ref().secret.to_bytes()[.. 32].as_ref(),
-  )
-  .unwrap();
-
-  assert_eq!(Ristretto::generator() * secret_key, public_key);
-  let threshold_keys =
-    musig::<Ristretto>(&musig_context(set), &Zeroizing::new(secret_key), &[public_key]).unwrap();
-
-  let sig = frost::tests::sign_without_caching(
-    &mut OsRng,
-    frost::tests::algorithm_machines(
-      &mut OsRng,
-      &Schnorrkel::new(b"substrate"),
-      &HashMap::from([(threshold_keys.params().i(), threshold_keys.into())]),
-    ),
-    &oraclize_values_message(&set, values),
-  );
-
-  // oraclize values
-  let _ =
-    publish_tx(serai, &SeraiGenesisLiquidity::oraclize_values(*values, Signature(sig.to_bytes())))
-      .await;
-}

substrate/client/tests/validator_sets.rs

@@ -6,7 +6,9 @@ use sp_core::{
 };
 
 use serai_client::{
-  primitives::{NETWORKS, NetworkId, BlockHash, insecure_pair_from_name},
+  primitives::{
+    NETWORKS, NetworkId, BlockHash, insecure_pair_from_name, FAST_EPOCH_DURATION, TARGET_BLOCK_TIME,
+  },
   validator_sets::{
     primitives::{Session, ValidatorSet, KeyPair},
     ValidatorSetsEvent,
@@ -326,22 +328,25 @@ async fn verify_session_and_active_validators(
   session: u32,
   participants: &[Public],
 ) {
-  // wait until the active session. This wait should be max 30 secs since the epoch time.
-  let block = tokio::time::timeout(core::time::Duration::from_secs(2 * 60), async move {
-    loop {
-      let mut block = serai.latest_finalized_block_hash().await.unwrap();
-      if session_for_block(serai, block, network).await < session {
-        // Sleep a block
-        tokio::time::sleep(core::time::Duration::from_secs(6)).await;
-        continue;
+  // wait until the active session.
+  let block = tokio::time::timeout(
+    core::time::Duration::from_secs(FAST_EPOCH_DURATION * TARGET_BLOCK_TIME * 2),
+    async move {
+      loop {
+        let mut block = serai.latest_finalized_block_hash().await.unwrap();
+        if session_for_block(serai, block, network).await < session {
+          // Sleep a block
+          tokio::time::sleep(core::time::Duration::from_secs(TARGET_BLOCK_TIME)).await;
+          continue;
+        }
+        while session_for_block(serai, block, network).await > session {
+          block = serai.block(block).await.unwrap().unwrap().header.parent_hash.0;
+        }
+        assert_eq!(session_for_block(serai, block, network).await, session);
+        break block;
       }
-      while session_for_block(serai, block, network).await > session {
-        block = serai.block(block).await.unwrap().unwrap().header.parent_hash.0;
-      }
-      assert_eq!(session_for_block(serai, block, network).await, session);
-      break block;
-    }
-  })
+    },
+  )
   .await
   .unwrap();
   let serai_for_block = serai.as_of(block);
@@ -358,10 +363,10 @@ async fn verify_session_and_active_validators(
 
   // make sure finalization continues as usual after the changes
   let current_finalized_block = serai.latest_finalized_block().await.unwrap().header.number;
-  tokio::time::timeout(core::time::Duration::from_secs(60), async move {
+  tokio::time::timeout(core::time::Duration::from_secs(TARGET_BLOCK_TIME * 10), async move {
     let mut finalized_block = serai.latest_finalized_block().await.unwrap().header.number;
     while finalized_block <= current_finalized_block + 2 {
-      tokio::time::sleep(core::time::Duration::from_secs(6)).await;
+      tokio::time::sleep(core::time::Duration::from_secs(TARGET_BLOCK_TIME)).await;
       finalized_block = serai.latest_finalized_block().await.unwrap().header.number;
     }
   })