Fix tests

knack/cli.py

@@ -22,6 +22,8 @@
 
 logger = get_logger(__name__)
 
+KNACK_TEST_FORCE_ENABLE_COLOR = False
+
 
 class CLI(object):  # pylint: disable=too-many-instance-attributes
     """ The main driver for the CLI """
@@ -100,8 +102,9 @@ def __init__(self,
 
         # Color is only enabled when all scenarios are met:
         #     1. [core] no_color config is not set
-        #     2. out_file (stdout by default) is a tty
+        #     2. stdout is a tty
         #     3. stderr is a tty
+        #     4. out_file is stdout
         # Otherwise, if the downstream command doesn't support color, Knack will fail with
         # BrokenPipeError: [Errno 32] Broken pipe, like `az --version | head --lines=1`
         # https://github.com/Azure/azure-cli/issues/13413
@@ -218,7 +221,7 @@ def invoke(self, args, initial_invocation_data=None, out_file=None):
         exit_code = 0
         try:
             out_file = out_file or self.out_file
-            if out_file is sys.stdout and self.enable_color:
+            if out_file is sys.stdout and self.enable_color or KNACK_TEST_FORCE_ENABLE_COLOR:
                 import colorama
                 colorama.init()
                 # point out_file to the new sys.stdout which is overwritten by colorama
@@ -260,7 +263,7 @@ def invoke(self, args, initial_invocation_data=None, out_file=None):
         finally:
             self.raise_event(EVENT_CLI_POST_EXECUTE)
 
-            if self.enable_color:
+            if self.enable_color or KNACK_TEST_FORCE_ENABLE_COLOR:
                 import colorama
                 colorama.deinit()
 

tests/test_help.py

@@ -5,32 +5,13 @@
 
 import sys
 import unittest
-import mock
-from six import StringIO
-
 
-from knack.help import ArgumentGroupRegistry, HelpObject
+import mock
 from knack.arguments import ArgumentsContext
-from knack.commands import CLICommand, CLICommandsLoader, CommandGroup
+from knack.commands import CLICommandsLoader, CommandGroup
 from knack.events import EVENT_PARSER_GLOBAL_CREATE
-from knack.invocation import CommandInvoker
-
-from tests.util import MockContext, DummyCLI
-
-io = {}
-
-
-def redirect_io(func):
-    def wrapper(self):
-        global io
-        old_stdout = sys.stdout
-        old_stderr = sys.stderr
-        sys.stdout = sys.stderr = io = StringIO()
-        func(self)
-        io.close()
-        sys.stdout = old_stdout
-        sys.stderr = old_stderr
-    return wrapper
+from knack.help import ArgumentGroupRegistry, HelpObject
+from tests.util import DummyCLI, redirect_io
 
 
 def example_handler(arg1, arg2=None, arg3=None):
@@ -206,7 +187,7 @@ def test_choice_list_with_ints(self):
         """ Ensure choice_list works with integer lists. """
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n1 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = 'Allowed values: 1, 2, 3'
         self.assertIn(expected, actual)
 
@@ -227,7 +208,7 @@ def test_help_long_and_short_description_from_docstring(self):
 
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n1 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = '\nCommand\n    {} n1 : Short summary here.\n        Long summary here. Still long summary.'.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -238,7 +219,7 @@ def test_help_long_and_short_description_from_yaml(self):
 
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n2 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = '\nCommand\n    {} n2 : YAML short summary.\n        YAML long summary. More summary.'.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -248,7 +229,7 @@ def test_help_long_description_multi_line(self):
 
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n3 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = '\nCommand\n    {} n3 : Short summary here.\n        Line1\n        line2.\n'.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -269,7 +250,7 @@ def test_help_params_documentations(self):
         Paragraph(s).
     --foobar3            : The foobar3.
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -307,7 +288,7 @@ def test_help_full_documentations(self):
         example details
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -338,7 +319,7 @@ def test_help_with_param_specified(self):
     --verbose          : Increase logging verbosity. Use --debug for full debug logs.
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 
@@ -358,7 +339,7 @@ def test_help_group_children(self):
     beta
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 
@@ -377,7 +358,7 @@ def test_help_missing_params(self):
             with self.assertRaises(SystemExit):
                 self.cli_ctx.invoke('n1 -a 1 --arg 2'.split())
 
-            actual = io.getvalue()
+            actual = self.io.getvalue()
             self.assertTrue('required' in actual and '-b' in actual)
 
     @redirect_io
@@ -395,7 +376,7 @@ def test_help_extra_params(self):
             with self.assertRaises(SystemExit):
                 self.cli_ctx.invoke('n1 -a 1 -b c -c extra'.split())
 
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = 'unrecognized arguments: -c extra'
         self.assertIn(expected, actual)
 
@@ -415,7 +396,7 @@ def test_help_group_help(self):
     n1 : This module does xyz one-line or so.
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 
@@ -455,7 +436,7 @@ def register_globals(_, **kwargs):
     --verbose          : Increase logging verbosity. Use --debug for full debug logs.
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = s.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 

tests/util.py

@@ -33,9 +33,10 @@ def wrapper(self):
         cli_logger.handlers.clear()
 
         sys.stdout = sys.stderr = self.io = StringIO()
-        func(self)
+        with mock.patch("knack.cli.KNACK_TEST_FORCE_ENABLE_COLOR", True):
+            func(self)
         self.io.close()
-        sys.stdout = original_stderr
+        sys.stdout = original_stdout
         sys.stderr = original_stderr
 
         # Remove the handlers added by CLI, so that the next invoke call init them again with the new stderr