diff options
Diffstat (limited to 'python')
| -rw-r--r-- | python/samba/tests/samba_tool/visualize.py | 73 |
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 |