pytest samba-tool visualize: extend colour tests for $NO_COLOR

As described at https://no-color.org/, the NO_COLOR environment
variable is a widely used defacto-ish standard for asking for no
colour. If someone goes

                NO_COLOR=whatever samba-tool ...

we want to assume they want no ANSI colour codes, as if they had used
--color=no. But first we want to test that, so here we are.

Signed-off-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>
Reviewed-by: Joseph Sutton <josephsutton@catalyst.net.nz>

1 files changed, 73 insertions, 0 deletions

diff --git a/python/samba/tests/samba_tool/visualize.py b/python/samba/tests/samba_tool/visualize.py
index 52fffc06fd4..b4385d0b716 100644
--- a/python/samba/tests/samba_tool/visualize.py
+++ b/python/samba/tests/samba_tool/visualize.py
@@ -28,6 +28,7 @@ import samba
 import os
 import tempfile
 import re
+from io import StringIO
 from samba.tests.samba_tool.base import SambaToolCmdTest
 from samba.kcc import ldif_import_export
 from samba.graph import COLOUR_SETS
@@ -58,6 +59,13 @@ MULTISITE_LDIF_DSAS = [
 ]
 
 
+class StringIOThinksItIsATTY(StringIO):
+    """A StringIO that claims to be a TTY for testing --color=auto,
+    by switching the stringIO class attribute."""
+    def isatty(self):
+        return True
+
+
 def samdb_from_ldif(ldif, tempdir, lp, dsa=None, tag=''):
     if dsa is None:
         dsa_name = 'default-DSA'
@@ -123,6 +131,71 @@ class SambaToolVisualizeLdif(SambaToolCmdTest):
 
             self.assertStringsEqual(monochrome, uncoloured, strip=True)
 
+    def assert_colour(self, text, has_colour=True, monochrome=None):
+        colour_re = re.compile('\033' r'\[[\d;]+m')
+        found = colour_re.search(text)
+        if has_colour:
+            self.assertTrue(found, text)
+        else:
+            self.assertFalse(found, text)
+        if monochrome is not None:
+            uncoloured = colour_re.sub('', text)
+            self.assertStringsEqual(monochrome, uncoloured, strip=True)
+
+    def test_colour_auto_tty(self):
+        """Assert the behaviour of --colour=auto with and without
+        NO_COLOUR on a fake tty"""
+        result, monochrome, err = self.runsubcmd("visualize", "ntdsconn",
+                                                 '-H', self.dburl,
+                                                 '--color=no', '-S')
+        self.assertCmdSuccess(result, monochrome, err)
+        self.assert_colour(monochrome, False)
+
+        try:
+            self.stringIO = StringIOThinksItIsATTY
+            old_no_color = os.environ.pop('NO_COLOR', None)
+            # First with no NO_COLOR env var. There should be colour.
+            result, out, err = self.runsubcmd("visualize", "ntdsconn",
+                                              '-H', self.dburl,
+                                              '-S',
+                                              '--color=auto')
+            self.assertCmdSuccess(result, out, err)
+            self.assert_colour(out, True, monochrome)
+
+            for env, opt, is_colour in [
+                    # NO_COLOR='' should be as if no NO_COLOR
+                    ['', '--color=auto', True],
+                    # NO_COLOR='1': we expect no colour
+                    ['1', '--color=auto', False],
+                    # NO_COLOR='no': we still expect no colour
+                    ['no', '--color=auto', False],
+                    # NO_COLOR=' ', alias for 'auto'
+                    ]:
+                os.environ['NO_COLOR'] = env
+                print(f" {env}, {opt}, {is_colour}")
+                result, out, err = self.runsubcmd("visualize", "ntdsconn",
+                                                  '-H', self.dburl,
+                                                  '-S',
+                                                  opt)
+                self.assertCmdSuccess(result, out, err)
+                self.assert_colour(out, is_colour, monochrome)
+
+                # with "-o -" output filename alias for stdout.
+                result, out, err = self.runsubcmd("visualize", "ntdsconn",
+                                                  '-H', self.dburl,
+                                                  '-S',
+                                                  opt,
+                                                  '-o', '-')
+                self.assertCmdSuccess(result, out, err)
+                self.assert_colour(out, is_colour, monochrome)
+
+        finally:
+            self.stringIO = StringIO
+            if old_no_color is None:
+                os.environ.pop('NO_COLOR', None)
+            else:
+                os.environ['NO_COLOR'] = old_no_color
+
     def test_import_ldif_xdot(self):
         """We can't test actual xdot, but using the environment we can
         persuade samba-tool that a script we write is xdot and ensure