From db74e6ecfecdc19b3821a908721b7c00b0808b22 Mon Sep 17 00:00:00 2001
From: 0xtekgrinder <72015889+0xtekgrinder@users.noreply.github.com>
Date: Sat, 20 Apr 2024 19:10:30 -0400
Subject: [PATCH] feat: add envExists cheatcode to check if a environment
 variable exists (#7744)

---
 crates/cheatcodes/assets/cheatcodes.json | 20 ++++++++++++++++++++
 crates/cheatcodes/spec/src/vm.rs         |  4 ++++
 crates/cheatcodes/src/env.rs             |  7 +++++++
 testdata/cheats/Vm.sol                   |  1 +
 testdata/default/cheats/Env.t.sol        |  8 ++++++++
 5 files changed, 40 insertions(+)

diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json
index 30a63009a59a..9ce73e931abe 100644
--- a/crates/cheatcodes/assets/cheatcodes.json
+++ b/crates/cheatcodes/assets/cheatcodes.json
@@ -3811,6 +3811,26 @@
       "status": "stable",
       "safety": "safe"
     },
+    {
+      "func": {
+        "id": "envExists",
+        "description": "Gets the environment variable `name` and returns true if it exists, else returns false.",
+        "declaration": "function envExists(string calldata name) external view returns (bool exists);",
+        "visibility": "external",
+        "mutability": "view",
+        "signature": "envExists(string)",
+        "selector": "0xce8365f9",
+        "selectorBytes": [
+          206,
+          131,
+          101,
+          249
+        ]
+      },
+      "group": "environment",
+      "status": "stable",
+      "safety": "safe"
+    },
     {
       "func": {
         "id": "envInt_0",
diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs
index e9d5a49caa51..39f0fd8b0006 100644
--- a/crates/cheatcodes/spec/src/vm.rs
+++ b/crates/cheatcodes/spec/src/vm.rs
@@ -1501,6 +1501,10 @@ interface Vm {
     #[cheatcode(group = Environment)]
     function setEnv(string calldata name, string calldata value) external;
 
+    /// Gets the environment variable `name` and returns true if it exists, else returns false.
+    #[cheatcode(group = Environment)]
+    function envExists(string calldata name) external view returns (bool exists);
+
     /// Gets the environment variable `name` and parses it as `bool`.
     /// Reverts if the variable was not found or could not be parsed.
     #[cheatcode(group = Environment)]
diff --git a/crates/cheatcodes/src/env.rs b/crates/cheatcodes/src/env.rs
index d9022d1b3f08..b8245dda1d00 100644
--- a/crates/cheatcodes/src/env.rs
+++ b/crates/cheatcodes/src/env.rs
@@ -26,6 +26,13 @@ impl Cheatcode for setEnvCall {
     }
 }
 
+impl Cheatcode for envExistsCall {
+    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+        let Self { name } = self;
+        Ok(env::var(name).is_ok().abi_encode())
+    }
+}
+
 impl Cheatcode for envBool_0Call {
     fn apply(&self, _state: &mut Cheatcodes) -> Result {
         let Self { name } = self;
diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol
index b04a07e1d54d..02eeac480e33 100644
--- a/testdata/cheats/Vm.sol
+++ b/testdata/cheats/Vm.sol
@@ -186,6 +186,7 @@ interface Vm {
     function envBytes32(string calldata name, string calldata delim) external view returns (bytes32[] memory value);
     function envBytes(string calldata name) external view returns (bytes memory value);
     function envBytes(string calldata name, string calldata delim) external view returns (bytes[] memory value);
+    function envExists(string calldata name) external view returns (bool exists);
     function envInt(string calldata name) external view returns (int256 value);
     function envInt(string calldata name, string calldata delim) external view returns (int256[] memory value);
     function envOr(string calldata name, bool defaultValue) external view returns (bool value);
diff --git a/testdata/default/cheats/Env.t.sol b/testdata/default/cheats/Env.t.sol
index 523ab34dc1f5..e325df2fa442 100644
--- a/testdata/default/cheats/Env.t.sol
+++ b/testdata/default/cheats/Env.t.sol
@@ -13,6 +13,14 @@ contract EnvTest is DSTest {
         vm.setEnv(key, val);
     }
 
+    function testEnvExists() public {
+        string memory key = "_foundryCheatcodeEnvExistsTestKey";
+        string memory val = "_foundryCheatcodeEnvExistsTestVal";
+        vm.setEnv(key, val);
+        require(vm.envExists(key), "envExists failed");
+        require(!vm.envExists("nonexistent"), "envExists failed");
+    }
+
     uint256 constant numEnvBoolTests = 2;
 
     function testEnvBool() public {