utils.astring: A couple of strip_console_codes fixes

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

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 <[email protected]>
+
+# 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()