From a757f20f38178fd4ce0d936688777c1f454544b0 Mon Sep 17 00:00:00 2001
From: Balaji Arun <balajia@vt.edu>
Date: Tue, 24 Oct 2023 14:02:54 -0700
Subject: [PATCH] [forge] Ability to set onchain config in tests

---
 crates/aptos-genesis/src/config.rs            | 5 +++++
 crates/aptos-genesis/src/lib.rs               | 2 +-
 crates/aptos/src/genesis/mod.rs               | 3 ++-
 terraform/helm/genesis/templates/genesis.yaml | 3 +++
 terraform/helm/genesis/values.yaml            | 2 ++
 testsuite/forge-cli/src/main.rs               | 7 ++++++-
 6 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/crates/aptos-genesis/src/config.rs b/crates/aptos-genesis/src/config.rs
index 26c49e9fb389c9..d4396f88a57b81 100644
--- a/crates/aptos-genesis/src/config.rs
+++ b/crates/aptos-genesis/src/config.rs
@@ -7,6 +7,7 @@ use aptos_types::{
     account_address::{AccountAddress, AccountAddressWithChecks},
     chain_id::ChainId,
     network_address::{DnsName, NetworkAddress, Protocol},
+    on_chain_config::OnChainConsensusConfig,
     transaction::authenticator::AuthenticationKey,
 };
 use aptos_vm_genesis::{AccountBalance, EmployeePool, Validator, ValidatorWithCommissionRate};
@@ -66,6 +67,9 @@ pub struct Layout {
     pub employee_vesting_start: Option<u64>,
     /// Duration of each vesting period (in seconds).
     pub employee_vesting_period_duration: Option<u64>,
+    /// Onchain Consensus Config
+    #[serde(default)]
+    pub on_chain_consensus_config: OnChainConsensusConfig,
 }
 
 impl Layout {
@@ -103,6 +107,7 @@ impl Default for Layout {
             total_supply: None,
             employee_vesting_start: Some(1663456089),
             employee_vesting_period_duration: Some(5 * 60), // 5 minutes
+            on_chain_consensus_config: OnChainConsensusConfig::default(),
         }
     }
 }
diff --git a/crates/aptos-genesis/src/lib.rs b/crates/aptos-genesis/src/lib.rs
index d0b2193f79c093..c11aef369c2e39 100644
--- a/crates/aptos-genesis/src/lib.rs
+++ b/crates/aptos-genesis/src/lib.rs
@@ -32,7 +32,7 @@ use aptos_vm_genesis::Validator;
 use std::convert::TryInto;
 
 /// Holder object for all pieces needed to generate a genesis transaction
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct GenesisInfo {
     /// ChainId for identifying the network
     chain_id: ChainId,
diff --git a/crates/aptos/src/genesis/mod.rs b/crates/aptos/src/genesis/mod.rs
index 5b8ff77a4e4f2a..07c042602c61f0 100644
--- a/crates/aptos/src/genesis/mod.rs
+++ b/crates/aptos/src/genesis/mod.rs
@@ -120,6 +120,7 @@ impl CliCommand<Vec<PathBuf>> for GenerateGenesis {
             (genesis_bytes, mainnet_genesis.generate_waypoint()?)
         } else {
             let mut test_genesis = fetch_genesis_info(self.git_options)?;
+            info!("Using Genesis Info: {:?}", test_genesis);
             let genesis_bytes = bcs::to_bytes(test_genesis.clone().get_genesis())
                 .map_err(|e| CliError::BCS(GENESIS_FILE, e))?;
             (genesis_bytes, test_genesis.generate_waypoint()?)
@@ -294,7 +295,7 @@ pub fn fetch_genesis_info(git_options: GitOptions) -> CliTypedResult<GenesisInfo
             voting_power_increase_limit: layout.voting_power_increase_limit,
             employee_vesting_start: layout.employee_vesting_start,
             employee_vesting_period_duration: layout.employee_vesting_period_duration,
-            consensus_config: OnChainConsensusConfig::default(),
+            consensus_config: layout.on_chain_consensus_config,
             execution_config: OnChainExecutionConfig::default_for_genesis(),
             gas_schedule: default_gas_schedule(),
         },
diff --git a/terraform/helm/genesis/templates/genesis.yaml b/terraform/helm/genesis/templates/genesis.yaml
index 6d5716a25e6fee..9aa1099d1afe73 100644
--- a/terraform/helm/genesis/templates/genesis.yaml
+++ b/terraform/helm/genesis/templates/genesis.yaml
@@ -25,6 +25,9 @@ data:
     rewards_apy_percentage: {{ .Values.chain.rewards_apy_percentage | int }}
     voting_duration_secs: {{ .Values.chain.voting_duration_secs | int }}
     voting_power_increase_limit: {{ .Values.chain.voting_power_increase_limit | int }}
+    {{- if len .Values.chain.on_chain_consensus_config > 1}}
+    on_chain_consensus_config: {{ .Values.chain.on_chain_consensus_config | mustToJson | quote }}
+    {{- end}}
 
 ---
 
diff --git a/terraform/helm/genesis/values.yaml b/terraform/helm/genesis/values.yaml
index 4a8c859e618594..dc5a7789e65fe9 100644
--- a/terraform/helm/genesis/values.yaml
+++ b/terraform/helm/genesis/values.yaml
@@ -31,6 +31,8 @@ chain:
   rewards_apy_percentage: 10
   # -- Minimum price per gas unit
   min_price_per_gas_unit: 1
+  # -- Onchain Consensus Config
+  on_chain_consensus_config:
 
 # -- Default image tag to use for all tools images
 imageTag: testnet
diff --git a/testsuite/forge-cli/src/main.rs b/testsuite/forge-cli/src/main.rs
index b43db6e4330c6a..a444ada632ee7f 100644
--- a/testsuite/forge-cli/src/main.rs
+++ b/testsuite/forge-cli/src/main.rs
@@ -20,7 +20,10 @@ use aptos_forge::{
 };
 use aptos_logger::{info, Level};
 use aptos_rest_client::Client as RestClient;
-use aptos_sdk::{move_types::account_address::AccountAddress, transaction_builder::aptos_stdlib};
+use aptos_sdk::{
+    move_types::account_address::AccountAddress, transaction_builder::aptos_stdlib,
+    types::on_chain_config::OnChainConsensusConfig,
+};
 use aptos_testcases::{
     compatibility_test::SimpleValidatorUpgrade,
     consensus_reliability_tests::ChangingWorkingQuorumTest,
@@ -1775,6 +1778,8 @@ fn realistic_env_max_load_test(
             // Have single epoch change in land blocking, and a few on long-running
             helm_values["chain"]["epoch_duration_secs"] =
                 (if long_running { 600 } else { 300 }).into();
+            helm_values["chain"]["on_chain_consensus_config"] =
+                OnChainConsensusConfig::default().into();
         }))
         // First start higher gas-fee traffic, to not cause issues with TxnEmitter setup - account creation
         .with_emit_job(