fix incorrect mocking way to unblock (#262)

* fix * fix style
Azure · Oct 9, 2020 · 9b59e4e · 9b59e4e
1 parent 75a1a79
commit 9b59e4e
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/azdev/operations/tests/test_benchmark.py b/azdev/operations/tests/test_benchmark.py
@@ -81,14 +81,18 @@ def test_load_all_commands_ok(self):
     def test_load_all_commands_fail(self):
         import sys
 
-        sys.modules["azure.cli.core"] = mock.MagicMock(side_effect=ImportError)
+        original_azure_cli_core_mod = sys.modules.get("azure.cli.core")
+        sys.modules["azure.cli.core"] = None
 
         with self.assertRaisesRegex(CLIError, "Azure CLI is not installed"):
             _benchmark_load_all_commands()
 
-        del sys.modules[
-            "azure.cli.core"
-        ]  # restore azure.cli.core to be unimported as the original
+        if original_azure_cli_core_mod:
+            sys.modules["azure.cli.core"] = original_azure_cli_core_mod
+        else:
+            del sys.modules[
+                "azure.cli.core"
+            ]  # restore azure.cli.core to the unimported status
 
 
 # class _MockedMapResultCounter: