diff --git a/codeflash/code_utils/shell_utils.py b/codeflash/code_utils/shell_utils.py index 30a5aadaa..79b211111 100644 --- a/codeflash/code_utils/shell_utils.py +++ b/codeflash/code_utils/shell_utils.py @@ -15,7 +15,9 @@ 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=" diff --git a/tests/test_shell_utils.py b/tests/test_shell_utils.py index 51d484064..c152ca7dd 100644 --- a/tests/test_shell_utils.py +++ b/tests/test_shell_utils.py @@ -64,6 +64,12 @@ def test_valid_api_key(self): ) as mock_file: self.assertEqual(read_api_key_from_shell_config(), None) 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") + @patch("codeflash.code_utils.shell_utils.get_shell_rc_path") def test_no_api_key(self, mock_get_shell_rc_path):