update to include oses

Author: KRRT7

tests/test_shell_utils.py

@@ -60,21 +60,33 @@ 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")
-            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")
+            
+            if os.name != "nt":
+                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")
+                
+                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")
+            
+            elif os.name == "nt":
+                with patch(
+                    "builtins.open", mock_open(read_data=f'REM set 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")
@@ -89,15 +101,27 @@ def test_no_api_key(self, mock_get_shell_rc_path):
     def test_malformed_api_key_export(self, mock_get_shell_rc_path):
         """Test with a malformed API key export."""
         mock_get_shell_rc_path.return_value = self.test_rc_path
-        with patch("builtins.open", mock_open(read_data=f"export API_KEY={self.api_key}\n")):
-            result = read_api_key_from_shell_config()
-            self.assertIsNone(result)
-        with patch("builtins.open", mock_open(read_data=f"CODEFLASH_API_KEY={self.api_key}\n")):
-            result = read_api_key_from_shell_config()
-            self.assertIsNone(result)
-        with patch("builtins.open", mock_open(read_data=f"export CODEFLASH_API_KEY=sk-{self.api_key}\n")):
-            result = read_api_key_from_shell_config()
-            self.assertIsNone(result)
+        
+        if os.name == "nt":
+            with patch("builtins.open", mock_open(read_data=f"set API_KEY={self.api_key}\n")):
+                result = read_api_key_from_shell_config()
+                self.assertIsNone(result)
+            with patch("builtins.open", mock_open(read_data=f"CODEFLASH_API_KEY={self.api_key}\n")):
+                result = read_api_key_from_shell_config()
+                self.assertIsNone(result)
+            with patch("builtins.open", mock_open(read_data=f"set CODEFLASH_API_KEY=sk-{self.api_key}\n")):
+                result = read_api_key_from_shell_config()
+                self.assertIsNone(result)
+        else:
+            with patch("builtins.open", mock_open(read_data=f"export API_KEY={self.api_key}\n")):
+                result = read_api_key_from_shell_config()
+                self.assertIsNone(result)
+            with patch("builtins.open", mock_open(read_data=f"CODEFLASH_API_KEY={self.api_key}\n")):
+                result = read_api_key_from_shell_config()
+                self.assertIsNone(result)
+            with patch("builtins.open", mock_open(read_data=f"export CODEFLASH_API_KEY=sk-{self.api_key}\n")):
+                result = read_api_key_from_shell_config()
+                self.assertIsNone(result)
 
     @patch("codeflash.code_utils.shell_utils.get_shell_rc_path")
     def test_multiple_api_key_exports(self, mock_get_shell_rc_path):
@@ -106,7 +130,7 @@ def test_multiple_api_key_exports(self, mock_get_shell_rc_path):
         if os.name == "nt":  # Windows
             first_export = 'set CODEFLASH_API_KEY=cf-firstkey'
             second_export = f'set CODEFLASH_API_KEY={self.api_key}'
-        else:  # Unix-like systems
+        else:
             first_export = 'export CODEFLASH_API_KEY="cf-firstkey"'
             second_export = f'export CODEFLASH_API_KEY="{self.api_key}"'
         with patch(