feat: add gas_snapshot_check flag to config, fix FORGE_SNAPSHOT_CHECK behavior

crates/config/src/lib.rs

@@ -200,6 +200,8 @@ pub struct Config {
     pub cache_path: PathBuf,
     /// where the gas snapshots are stored
     pub snapshots: PathBuf,
+    /// whether to check for differences against previously stored gas snapshots
+    pub gas_snapshot_check: bool,
     /// where the broadcast logs are stored
     pub broadcast: PathBuf,
     /// additional solc allow paths for `--allow-paths`
@@ -2316,6 +2318,7 @@ impl Default for Config {
             cache_path: "cache".into(),
             broadcast: "broadcast".into(),
             snapshots: "snapshots".into(),
+            gas_snapshot_check: false,
             allow_paths: vec![],
             include_paths: vec![],
             force: false,

crates/forge/bin/cmd/test/mod.rs

@@ -118,6 +118,10 @@ pub struct TestArgs {
     #[arg(long, env = "FORGE_GAS_REPORT")]
     gas_report: bool,
 
+    /// Check gas snapshots against previous runs.
+    #[arg(long, env = "FORGE_SNAPSHOT_CHECK")]
+    gas_snapshot_check: Option<bool>,
+
     /// Exit with code 0 even if a test fails.
     #[arg(long, env = "FORGE_ALLOW_FAILURE")]
     allow_failure: bool,
@@ -662,9 +666,25 @@ impl TestArgs {
 
             // Write gas snapshots to disk if any were collected.
             if !gas_snapshots.is_empty() {
-                // Check for differences in gas snapshots if `FORGE_SNAPSHOT_CHECK` is set.
+                // By default `gas_snapshot_check` is set to `false` in the config.
+                //
+                // The user can either:
+                // - Set `FORGE_SNAPSHOT_CHECK=true` in the environment.
+                // - Pass `--gas-snapshot-check=true` as a CLI argument.
+                // - Set `gas_snapshot_check = true` in the config.
+                //
+                // If the user passes `--gas-snapshot-check=<bool>` then it will override the config
+                // and the environment variable, disabling the check if `false` is passed.
+                //
                 // Exiting early with code 1 if differences are found.
-                if std::env::var("FORGE_SNAPSHOT_CHECK").is_ok() {
+                let should_check_for_differences =
+                    if let Some(gas_snapshot_check) = self.gas_snapshot_check {
+                        gas_snapshot_check
+                    } else {
+                        config.gas_snapshot_check
+                    };
+
+                if should_check_for_differences {
                     let differences_found = gas_snapshots.clone().into_iter().fold(
                         false,
                         |mut found, (group, snapshots)| {

crates/forge/tests/cli/config.rs

@@ -40,6 +40,7 @@ forgetest!(can_extract_config_values, |prj, cmd| {
         cache: true,
         cache_path: "test-cache".into(),
         snapshots: "snapshots".into(),
+        gas_snapshot_check: false,
         broadcast: "broadcast".into(),
         force: true,
         evm_version: EvmVersion::Byzantium,
@@ -978,6 +979,7 @@ libraries = []
 cache = true
 cache_path = "cache"
 snapshots = "snapshots"
+gas_snapshot_check = false
 broadcast = "broadcast"
 allow_paths = []
 include_paths = []
@@ -1132,6 +1134,7 @@ exclude = []
   "cache": true,
   "cache_path": "cache",
   "snapshots": "snapshots",
+  "gas_snapshot_check": false,
   "broadcast": "broadcast",
   "allow_paths": [],
   "include_paths": [],
@@ -1365,3 +1368,181 @@ optimizer_runs = 0
 
 "#]]);
 });
+
+forgetest_init!(test_gas_snapshot_check_config, |prj, cmd| {
+    // Default settings: gas_snapshot_check disabled.
+    cmd.forge_fuse().args(["config"]).assert_success().stdout_eq(str![[r#"
+...
+gas_snapshot_check = false
+...
+
+"#]]);
+
+    prj.insert_ds_test();
+
+    prj.add_source(
+        "Flare.sol",
+        r#"
+contract Flare {
+    bytes32[] public data;
+
+    function run(uint256 n_) public {
+        for (uint256 i = 0; i < n_; i++) {
+            data.push(keccak256(abi.encodePacked(i)));
+        }
+    }
+}
+    "#,
+    )
+    .unwrap();
+
+    let test_contract = |n: u32| {
+        format!(
+            r#"
+import "./test.sol";
+import "./Flare.sol";
+
+interface Vm {{
+    function startSnapshotGas(string memory name) external;
+    function stopSnapshotGas() external returns (uint256);
+}}
+
+contract GasSnapshotCheckTest is DSTest {{
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    Flare public flare;
+
+    function setUp() public {{
+        flare = new Flare();
+    }}
+
+    function testSnapshotGasSectionExternal() public {{
+        vm.startSnapshotGas("testAssertGasExternal");
+        flare.run({n});
+        vm.stopSnapshotGas();
+    }}
+}}
+        "#
+        )
+    };
+
+    // Assert that gas_snapshot_check is disabled by default
+    prj.add_source("GasSnapshotCheckTest.sol", &test_contract(1)).unwrap();
+    cmd.forge_fuse().args(["test"]).assert_success().stdout_eq(str![[r#"
+...
+Ran 1 test for src/GasSnapshotCheckTest.sol:GasSnapshotCheckTest
+[PASS] testSnapshotGasSectionExternal() ([GAS])
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+...
+"#]]);
+
+    // Enable gas_snapshot_check
+    let config = Config { gas_snapshot_check: true, ..Default::default() };
+    prj.write_config(config);
+    cmd.forge_fuse().args(["config"]).assert_success().stdout_eq(str![[r#"
+...
+gas_snapshot_check = true
+...
+
+"#]]);
+
+    // Replace the test contract with a new one that will fail the gas snapshot check
+    prj.add_source("GasSnapshotCheckTest.sol", &test_contract(2)).unwrap();
+    cmd.forge_fuse().args(["test"]).assert_failure().stderr_eq(str![[r#"
+...
+[GasSnapshotCheckTest] Failed to match snapshots:
+- [testAssertGasExternal] [..] → [..]
+
+Error: Snapshots differ from previous run
+...
+"#]]);
+
+    // Disable gas_snapshot_check, assert that running the test will pass
+    let config = Config { gas_snapshot_check: false, ..Default::default() };
+    prj.write_config(config);
+    cmd.forge_fuse().args(["test"]).assert_success().stdout_eq(str![[r#"
+...
+Ran 1 test for src/GasSnapshotCheckTest.sol:GasSnapshotCheckTest
+[PASS] testSnapshotGasSectionExternal() ([GAS])
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+...
+"#]]);
+
+    // Re-enable gas_snapshot_check
+    // Assert that the new value has been stored from the previous run and re-run the test
+    let config = Config { gas_snapshot_check: true, ..Default::default() };
+    prj.write_config(config);
+    cmd.forge_fuse().args(["test"]).assert_success().stdout_eq(str![[r#"
+...
+Ran 1 test for src/GasSnapshotCheckTest.sol:GasSnapshotCheckTest
+[PASS] testSnapshotGasSectionExternal() ([GAS])
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+...
+"#]]);
+
+    // Replace the test contract with a new one that will fail the gas_snapshot_check
+    prj.add_source("GasSnapshotCheckTest.sol", &test_contract(3)).unwrap();
+    cmd.forge_fuse().args(["test"]).assert_failure().stderr_eq(str![[r#"
+...
+[GasSnapshotCheckTest] Failed to match snapshots:
+- [testAssertGasExternal] [..] → [..]
+
+Error: Snapshots differ from previous run
+...
+"#]]);
+
+    // Test that `--gas-snapshot-check=false` flag can be used to disable the gas_snapshot_check
+    cmd.forge_fuse().args(["test", "--gas-snapshot-check=false"]).assert_success().stdout_eq(str![
+        [r#"
+...
+Ran 1 test for src/GasSnapshotCheckTest.sol:GasSnapshotCheckTest
+[PASS] testSnapshotGasSectionExternal() ([GAS])
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+...
+"#]
+    ]);
+
+    // Disable gas_snapshot_check in the config file
+    // Enable using `FORGE_SNAPSHOT_CHECK` environment variable
+    // Assert that this will override the config file value
+    let config = Config { gas_snapshot_check: false, ..Default::default() };
+    prj.write_config(config);
+    prj.add_source("GasSnapshotCheckTest.sol", &test_contract(4)).unwrap();
+    cmd.forge_fuse();
+    cmd.env("FORGE_SNAPSHOT_CHECK", "true");
+    cmd.args(["test"]).assert_failure().stderr_eq(str![[r#"
+...
+[GasSnapshotCheckTest] Failed to match snapshots:
+- [testAssertGasExternal] [..] → [..]
+
+Error: Snapshots differ from previous run
+...
+"#]]);
+
+    // Assert that `--gas-snapshot-check=true` flag can be used to enable the gas_snapshot_check
+    // even when `FORGE_SNAPSHOT_CHECK` is set to false in the environment variable
+    cmd.forge_fuse();
+    cmd.env("FORGE_SNAPSHOT_CHECK", "false");
+    cmd.args(["test", "--gas-snapshot-check=true"]).assert_failure().stderr_eq(str![[r#"
+...
+[GasSnapshotCheckTest] Failed to match snapshots:
+- [testAssertGasExternal] [..] → [..]
+
+Error: Snapshots differ from previous run
+...
+"#]]);
+
+    //
+
+    // Finally assert that `--gas-snapshot-check=false` flag can be used to disable the
+    // gas_snapshot_check even when `FORGE_SNAPSHOT_CHECK` is set to true
+    cmd.forge_fuse();
+    cmd.env("FORGE_SNAPSHOT_CHECK", "true");
+    cmd.args(["test", "--gas-snapshot-check=false"]).assert_success().stdout_eq(str![[r#"
+...
+Ran 1 test for src/GasSnapshotCheckTest.sol:GasSnapshotCheckTest
+[PASS] testSnapshotGasSectionExternal() ([GAS])
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+...
+"#]]);
+});