diff --git a/aexpect/utils/astring.py b/aexpect/utils/astring.py index 4e58313..34936ff 100644 --- a/aexpect/utils/astring.py +++ b/aexpect/utils/astring.py @@ -35,8 +35,9 @@ def strip_console_codes(output, custom_codes=None): return_str = "" index = 0 output = f"\x1b[m{output}" - console_codes = "%[G@8]|\\[[@A-HJ-MPXa-hl-nqrsu\\`]" - console_codes += "|\\[[\\d;]+[HJKgqnrm]|#8|\\([B0UK]|\\)|\\[\\?2004[lh]" + console_codes = "^%[G@8]|^\\[[@A-HJ-MPXa-hl-nqrsu\\`]" + console_codes += "|^\\[[\\d;]+[HJKgqnrm]|^#8|^\\([B0UK]|^\\)" + console_codes += "|^\\[\\?2004[lh]|^c|^\\[!p|^\\]\\d+" if custom_codes is not None and custom_codes not in console_codes: console_codes += f"|{custom_codes}" while index < len(output): @@ -54,7 +55,7 @@ def strip_console_codes(output, custom_codes=None): try: special_code = re.findall(console_codes, tmp_word)[0] except IndexError as error: - if index + tmp_index < len(output): + if index < len(output): raise ValueError(f"{tmp_word} is not included in the known " "console codes list " f"{console_codes}") from error diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..6e9c987 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,45 @@ +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation; either version 2 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. +# +# See LICENSE for more details. +# +# Copyright: Red Hat Inc. 2024 +# Author: Lukas Doktor + +# selftests pylint: disable=C0111,C0111 + +import unittest + +from aexpect.utils import astring + + +class Astring(unittest.TestCase): + + def test_strip_console_codes(self): + """ + Try various strip_console_codes + """ + strip = astring.strip_console_codes + self.assertEqual("simple color test", + strip("simple\x1b[33;1m color \x1b[0mtest")) + self.assertEqual("", + strip("\x1bskip-full-text")) + self.assertEqual("ignores last", + strip("ignores last\x1bbad")) + self.assertEqual("skips c [!p and ]104", + strip("skips\x1bc c \x1b[!p[!p and ]104\x1b]104")) + self.assertRaisesRegex(ValueError, "only is not included", strip, + "ignores\x1bonly\x1blast\x1bbad") + self.assertRaisesRegex(ValueError, "invalid-prefix.*included", strip, + "\x1binvalid-prefix[33;1mconsole code " + "must fail\x1b") + + +if __name__ == '__main__': + unittest.main()