utils.astring: Enforce the console code starts with the code

previously we only checked the console_code is found anywhere in the text between the two \x1b and then we blindly removed first N characters. Let's ensure we check the console code regexp from the beginning. Signed-off-by: Lukáš Doktor <[email protected]>
avocado-framework · Aug 19, 2024 · 506e44a · 506e44a
1 parent 10da7da
commit 506e44a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git 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]"
     if custom_codes is not None and custom_codes not in console_codes:
         console_codes += f"|{custom_codes}"
     while index < len(output):

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -34,6 +34,9 @@ def test_strip_console_codes(self):
                          strip("ignores last\x1bbad"))
         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__':