2 files changed, 96 insertions, 0 deletions

diff --git a/bitbake/bin/bitbake-selftest b/bitbake/bin/bitbake-selftest
index 041a2719f0..e84d6a559f 100755
--- a/bitbake/bin/bitbake-selftest
+++ b/bitbake/bin/bitbake-selftest
@@ -18,6 +18,7 @@ except RuntimeError as exc:
     sys.exit(str(exc))
 
 tests = ["bb.tests.codeparser",
+         "bb.tests.color",
          "bb.tests.cooker",
          "bb.tests.cow",
          "bb.tests.data",
diff --git a/bitbake/lib/bb/tests/color.py b/bitbake/lib/bb/tests/color.py
new file mode 100644
index 0000000000..bf03750c69
--- /dev/null
+++ b/bitbake/lib/bb/tests/color.py
@@ -0,0 +1,95 @@
+#
+# BitBake Test for ANSI color code filtering
+#
+# Copyright (C) 2020  Agilent Technologies, Inc.
+# Author: Chris Laplante <chris.laplante@agilent.com>
+#
+# SPDX-License-Identifier: MIT
+#
+
+import unittest
+import bb.progress
+import bb.data
+import bb.event
+from bb.progress import filter_color, filter_color_n
+import io
+import re
+
+
+class ProgressWatcher:
+    def __init__(self):
+        self._reports = []
+
+    def handle_event(self, event):
+        self._reports.append((event.progress, event.rate))
+
+    def reports(self):
+        return self._reports
+
+
+class ColorCodeTests(unittest.TestCase):
+    def setUp(self):
+        self.d = bb.data.init()
+        self._progress_watcher = ProgressWatcher()
+        bb.event.register("bb.build.TaskProgress", self._progress_watcher.handle_event)
+
+    def tearDown(self):
+        bb.event.remove("bb.build.TaskProgress", None)
+
+    def test_filter_color(self):
+        input_string = "~~~~~~~~~~~~^~~~~~~~"
+        filtered = filter_color(input_string)
+        self.assertEqual(filtered, "~~~~~~~~~~~~^~~~~~~~")
+
+    def test_filter_color_n(self):
+        input_string = "~~~~~~~~~~~~^~~~~~~~"
+        filtered, code_count = filter_color_n(input_string)
+        self.assertEqual(filtered, "~~~~~~~~~~~~^~~~~~~~")
+        self.assertEqual(code_count, 4)
+
+    def test_LineFilterProgressHandler_color_filtering(self):
+        class CustomProgressHandler(bb.progress.LineFilterProgressHandler):
+            PROGRESS_REGEX = re.compile(r"Progress: (?P<progress>\d+)%")
+
+            def writeline(self, line):
+                match = self.PROGRESS_REGEX.match(line)
+                if match:
+                    self.update(int(match.group("progress")))
+                    return False
+                return True
+
+        buffer = io.StringIO()
+        handler = CustomProgressHandler(self.d, buffer)
+        handler.write("Program output!\n")
+        handler.write("More output!\n")
+        handler.write("Progress: 10%\n") # 10%
+        handler.write("Even more\n")
+        handler.write("Progress: 50%\n") # 50%
+        handler.write("Progress: 60%\n") # 60%
+        handler.write("Progress: 100%\n") # 100%
+
+        expected = [(10, None), (50, None), (60, None), (100, None)]
+        self.assertEqual(self._progress_watcher.reports(), expected)
+
+        self.assertEqual(buffer.getvalue(), "Program output!\nMore output!\nEven more\n")
+
+    def test_BasicProgressHandler_color_filtering(self):
+        buffer = io.StringIO()
+        handler = bb.progress.BasicProgressHandler(self.d, outfile=buffer)
+        handler.write("1%\n") # 1%
+        handler.write("2%\n") # 2%
+        handler.write("10%\n") # 10%
+        handler.write("100%\n") # 100%
+
+        expected = [(0, None), (1, None), (2, None), (10, None), (100, None)]
+        self.assertListEqual(self._progress_watcher.reports(), expected)
+
+    def test_OutOfProgressHandler_color_filtering(self):
+        buffer = io.StringIO()
+        handler = bb.progress.OutOfProgressHandler(self.d, r'(\d+) of (\d+)', outfile=buffer)
+        handler.write("Text text 1 of 5") # 1/5
+        handler.write("Text text 3 of 5") # 3/5
+        handler.write("Text text 5 of 5") # 5/5
+
+        expected = [(0, None), (20.0, None), (60.0, None), (100.0, None)]
+        self.assertListEqual(self._progress_watcher.reports(), expected)