use decorator disable_color

knack/cli.py

@@ -191,9 +191,9 @@ def invoke(self, args, initial_invocation_data=None, out_file=None):
             if self.enable_color:
                 import colorama
                 colorama.init()
-            if self.out_file == sys.__stdout__:
-                # point out_file to the new sys.stdout which is overwritten by colorama
-                self.out_file = sys.stdout
+                if self.out_file == sys.__stdout__:
+                    # point out_file to the new sys.stdout which is overwritten by colorama
+                    self.out_file = sys.stdout
 
             args = self.completion.get_completion_args() or args
             out_file = out_file or self.out_file

tests/test_deprecation.py

@@ -15,7 +15,7 @@
 from knack.arguments import ArgumentsContext
 from knack.commands import CLICommand, CLICommandsLoader, CommandGroup
 
-from tests.util import DummyCLI, redirect_io
+from tests.util import DummyCLI, redirect_io, disable_color
 
 
 def example_handler(arg1, arg2=None, arg3=None):
@@ -140,17 +140,16 @@ def test_deprecate_command_expiring_execute_no_color(self):
     @redirect_io
     def test_deprecate_command_expired_execute(self):
         """ Ensure expired command cannot be reached. """
-        self.cli_ctx.enable_color = False
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('cmd5 -h'.split())
         actual = self.io.getvalue()
-        expected = """The most similar choices to 'cmd5'"""
+        expected = """cli: 'cmd5' is not in the 'cli' command group."""
         self.assertIn(expected, actual)
 
     @redirect_io
+    @disable_color
     def test_deprecate_command_expired_execute_no_color(self):
         """ Ensure error is displayed without color. """
-        self.cli_ctx.enable_color = False
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('cmd5 -h'.split())
         actual = self.io.getvalue()
@@ -249,9 +248,9 @@ def test_deprecate_command_group_help_expiring(self):
         self.assertIn(expected, actual)
 
     @redirect_io
+    @disable_color
     def test_deprecate_command_group_help_expiring_no_color(self):
         """ Ensure warning is displayed without color. """
-        self.cli_ctx.enable_color = False
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('group4 -h'.split())
         actual = self.io.getvalue()
@@ -447,9 +446,9 @@ def test_deprecate_options_execute_expiring(self):
         self.assertIn(expected, actual)
 
     @redirect_io
+    @disable_color
     def test_deprecate_options_execute_expiring_no_color(self):
         """ Ensure error is displayed without color. """
-        self.cli_ctx.enable_color = False
         self.cli_ctx.invoke('arg-test --arg1 foo --opt1 bar --alt4 bar'.split())
         actual = self.io.getvalue()
         expected = "WARNING: Option '--alt4' has been deprecated and will be removed in version '1.0.0'. Use '--opt4' instead."

tests/test_preview.py

@@ -18,7 +18,7 @@
 from knack.arguments import ArgumentsContext
 from knack.commands import CLICommandsLoader, CommandGroup
 
-from tests.util import DummyCLI, redirect_io
+from tests.util import DummyCLI, redirect_io, disable_color
 
 
 def example_handler(arg1, arg2=None, arg3=None):
@@ -89,9 +89,9 @@ def test_preview_command_plain_execute(self):
         self.assertIn(expected, actual)
 
     @redirect_io
+    @disable_color
     def test_preview_command_plain_execute_no_color(self):
         """ Ensure warning is displayed without color. """
-        self.cli_ctx.enable_color = False
         self.cli_ctx.invoke('cmd1 -b b'.split())
         actual = self.io.getvalue()
         self.assertIn("WARNING: This command is in preview. It may be changed/removed in a future release.", actual)
@@ -105,9 +105,9 @@ def test_preview_command_implicitly_execute(self):
         self.assertIn(expected, actual)
 
     @redirect_io
+    @disable_color
     def test_preview_command_implicitly_no_color(self):
         """ Ensure warning is displayed without color. """
-        self.cli_ctx.enable_color = False
         self.cli_ctx.invoke('grp1 cmd1 -b b'.split())
         actual = self.io.getvalue()
         expected = "WARNING: Command group 'grp1' is in preview. It may be changed/removed in a future release."
@@ -159,9 +159,9 @@ def test_preview_command_group_help_plain(self):
         self.assertEqual(expected, actual)
 
     @redirect_io
+    @disable_color
     def test_preview_command_group_help_plain_no_color(self):
         """ Ensure warning is displayed without color. """
-        self.cli_ctx.enable_color = False
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('group1 -h'.split())
         actual = self.io.getvalue()
@@ -234,9 +234,9 @@ def test_preview_arguments_command_help(self):
         self.assertIn(expected, actual)
 
     @redirect_io
+    @disable_color
     def test_preview_arguments_command_help_no_color(self):
         """ Ensure warning is displayed without color. """
-        self.cli_ctx.enable_color = False
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('arg-test -h'.split())
         actual = self.io.getvalue()
@@ -259,9 +259,9 @@ def test_preview_arguments_execute(self):
         self.assertIn(action_expected, actual)
 
     @redirect_io
+    @disable_color
     def test_preview_arguments_execute_no_color(self):
         """ Ensure warning is displayed without color. """
-        self.cli_ctx.enable_color = False
         self.cli_ctx.invoke('arg-test --arg1 foo --opt1 bar'.split())
         actual = self.io.getvalue()
         preview_expected = "WARNING: Argument '--arg1' is in preview. It may be changed/removed in a future release."

tests/util.py

@@ -12,6 +12,8 @@
 import shutil
 import os
 from six import StringIO
+import logging
+from knack.log import CLI_LOGGER_NAME
 
 from knack.cli import CLI, CLICommandsLoader, CommandInvoker
 
@@ -29,6 +31,21 @@ def wrapper(self):
         self.io.close()
         sys.stdout = original_stderr
         sys.stderr = original_stderr
+
+        # Remove the handlers added by CLI, so that the next invoke call init them again with the new stderr
+        # Otherwise, the handlers will write to a closed StringIO from a preview test
+        root_logger = logging.getLogger()
+        cli_logger = logging.getLogger(CLI_LOGGER_NAME)
+        root_logger.handlers = root_logger.handlers[:-1]
+        cli_logger.handlers = cli_logger.handlers[:-1]
+    return wrapper
+
+
+def disable_color(func):
+    def wrapper(self):
+        self.cli_ctx.enable_color = False
+        func(self)
+        self.cli_ctx.enable_color = True
     return wrapper