Merge pull request #462 from codeflash-ai/enhancement/better-shell-regex

[Enhancement] Better shell regex for matching api key export

Author: mohammedahmed18

codeflash/code_utils/shell_utils.py

@@ -15,7 +15,7 @@
     SHELL_RC_EXPORT_PATTERN = re.compile(r"^set CODEFLASH_API_KEY=(cf-.*)$", re.MULTILINE)
     SHELL_RC_EXPORT_PREFIX = "set CODEFLASH_API_KEY="
 else:
-    SHELL_RC_EXPORT_PATTERN = re.compile(r'^export CODEFLASH_API_KEY="?(cf-[^\s"]+)"?$', re.MULTILINE)
+    SHELL_RC_EXPORT_PATTERN = re.compile(r'^(?!#)export CODEFLASH_API_KEY=[\'"]?(cf-[^\s"]+)[\'"]$', re.MULTILINE)
     SHELL_RC_EXPORT_PREFIX = "export CODEFLASH_API_KEY="
 
 

tests/test_shell_utils.py

@@ -54,6 +54,16 @@ def test_valid_api_key(self):
             ) as mock_file:
                 self.assertEqual(read_api_key_from_shell_config(), self.api_key)
                 mock_file.assert_called_once_with(self.test_rc_path, encoding="utf8")
+            with patch(
+                "builtins.open", mock_open(read_data=f'export CODEFLASH_API_KEY=\'{self.api_key}\'\n')
+            ) as mock_file:
+                self.assertEqual(read_api_key_from_shell_config(), self.api_key)
+                mock_file.assert_called_once_with(self.test_rc_path, encoding="utf8")
+            with patch(
+                "builtins.open", mock_open(read_data=f'#export CODEFLASH_API_KEY=\'{self.api_key}\'\n')
+            ) as mock_file:
+                self.assertEqual(read_api_key_from_shell_config(), None)
+                mock_file.assert_called_once_with(self.test_rc_path, encoding="utf8")
 
     @patch("codeflash.code_utils.shell_utils.get_shell_rc_path")
     def test_no_api_key(self, mock_get_shell_rc_path):