Add CI-friendly progress output for tests (apache#24236)

This is the first step to run breeze tests in parallel in CI.
This flag adds "limited progress" output when running tests
which means that the runnig tests will just print few lines with
percent progress and color status indication from last few
progress lines of Pytest output, but when it completes, the whole output is
printed in a CI group - colored depending on status.

The final version (wnen we implement parallel test execution) should
also defer writing the output to until all tests are completed, but
this should be a follow-up PR.

TESTING.rst

@@ -182,6 +182,21 @@ You can also specify individual tests or a group of tests:
 
     breeze tests --db-reset tests/core/test_core.py::TestCore
 
+You can also limit the tests to execute to specific group of tests
+
+.. code-block:: bash
+
+    breeze tests --test-type Core
+
+
+You can also write tests in "limited progress" mode (useful in the future to run CI). In this mode each
+test just prints "percentage" summary of the run as single line and only dumps full output of the test
+after it completes.
+
+.. code-block:: bash
+
+    breeze tests --test-type Core --limit-progress-output
+
 
 Running Tests of a specified type from the Host
 -----------------------------------------------

dev/breeze/src/airflow_breeze/commands/configuration_and_maintenance_commands.py

@@ -157,9 +157,9 @@ def cleanup(verbose: bool, dry_run: bool, github_repository: str, all: bool, ans
         )
         images = command_result.stdout.splitlines() if command_result and command_result.stdout else []
         if images:
-            get_console().print("[light_blue]Removing images:[/]")
+            get_console().print("[info]Removing images:[/]")
             for image in images:
-                get_console().print(f"[light_blue] * {image}[/]")
+                get_console().print(f"[info] * {image}[/]")
             get_console().print()
             docker_rmi_command_to_execute = [
                 'docker',
@@ -173,7 +173,7 @@ def cleanup(verbose: bool, dry_run: bool, github_repository: str, all: bool, ans
             elif given_answer == Answer.QUIT:
                 sys.exit(0)
         else:
-            get_console().print("[light_blue]No locally downloaded images to remove[/]\n")
+            get_console().print("[info]No locally downloaded images to remove[/]\n")
     get_console().print("Pruning docker images")
     given_answer = user_confirm("Are you sure with the removal?")
     if given_answer == Answer.YES:

dev/breeze/src/airflow_breeze/commands/testing_commands.py

@@ -14,35 +14,46 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
+import errno
 import os
+import re
+import shutil
+import subprocess
 import sys
-from typing import Tuple
+import tempfile
+from threading import Event, Thread
+from time import sleep
+from typing import Dict, List, Tuple
 
 import click
 
 from airflow_breeze.commands.main_command import main
 from airflow_breeze.global_constants import ALLOWED_TEST_TYPES
 from airflow_breeze.params.build_prod_params import BuildProdParams
 from airflow_breeze.params.shell_params import ShellParams
+from airflow_breeze.utils.ci_group import ci_group
 from airflow_breeze.utils.common_options import (
+    option_backend,
     option_db_reset,
     option_dry_run,
     option_github_repository,
     option_image_name,
     option_image_tag,
     option_integration,
+    option_mssql_version,
+    option_mysql_version,
+    option_postgres_version,
     option_python,
     option_verbose,
 )
-from airflow_breeze.utils.console import get_console
+from airflow_breeze.utils.console import get_console, message_type_from_return_code
 from airflow_breeze.utils.custom_param_types import BetterChoice
 from airflow_breeze.utils.docker_command_utils import (
     get_env_variables_for_docker_commands,
     perform_environment_checks,
 )
 from airflow_breeze.utils.run_tests import run_docker_compose_tests
-from airflow_breeze.utils.run_utils import run_command
+from airflow_breeze.utils.run_utils import RunCommandResult, run_command
 
 TESTING_COMMANDS = {
     "name": "Testing",
@@ -55,8 +66,8 @@
             "name": "Docker-compose tests flag",
             "options": [
                 "--image-name",
-                "--python",
                 "--image-tag",
+                "--python",
             ],
         }
     ],
@@ -66,7 +77,13 @@
             "options": [
                 "--integration",
                 "--test-type",
+                "--limit-progress-output",
                 "--db-reset",
+                "--backend",
+                "--python",
+                "--postgres-version",
+                "--mysql-version",
+                "--mssql-version",
             ],
         }
     ],
@@ -112,6 +129,91 @@ def docker_compose_tests(
     sys.exit(return_code)
 
 
+class MonitoringThread(Thread):
+    """Thread class with a stop() method. The thread itself has to check
+    regularly for the stopped() condition."""
+
+    def __init__(self, title: str, file_name: str):
+        super().__init__(target=self.peek_percent_at_last_lines_of_file, daemon=True)
+        self._stop_event = Event()
+        self.title = title
+        self.file_name = file_name
+
+    def peek_percent_at_last_lines_of_file(self) -> None:
+        max_line_length = 400
+        matcher = re.compile(r"^.*\[([^\]]*)\]$")
+        while not self.stopped():
+            if os.path.exists(self.file_name):
+                try:
+                    with open(self.file_name, 'rb') as temp_f:
+                        temp_f.seek(-(max_line_length * 2), os.SEEK_END)
+                        tail = temp_f.read().decode()
+                    try:
+                        two_last_lines = tail.splitlines()[-2:]
+                        previous_no_ansi_line = escape_ansi(two_last_lines[0])
+                        m = matcher.match(previous_no_ansi_line)
+                        if m:
+                            get_console().print(f"[info]{self.title}:[/] {m.group(1).strip()}")
+                            print(f"\r{two_last_lines[0]}\r")
+                            print(f"\r{two_last_lines[1]}\r")
+                    except IndexError:
+                        pass
+                except OSError as e:
+                    if e.errno == errno.EINVAL:
+                        pass
+                    else:
+                        raise
+            sleep(5)
+
+    def stop(self):
+        self._stop_event.set()
+
+    def stopped(self):
+        return self._stop_event.is_set()
+
+
+def escape_ansi(line):
+    ansi_escape = re.compile(r'(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]')
+    return ansi_escape.sub('', line)
+
+
+def run_with_progress(
+    cmd: List[str],
+    env_variables: Dict[str, str],
+    test_type: str,
+    python: str,
+    backend: str,
+    version: str,
+    verbose: bool,
+    dry_run: bool,
+) -> RunCommandResult:
+    title = f"Running tests: {test_type}, Python: {python}, Backend: {backend}:{version}"
+    try:
+        with tempfile.NamedTemporaryFile(mode='w+t', delete=False) as f:
+            get_console().print(f"[info]Starting test = {title}[/]")
+            thread = MonitoringThread(title=title, file_name=f.name)
+            thread.start()
+            try:
+                result = run_command(
+                    cmd,
+                    verbose=verbose,
+                    dry_run=dry_run,
+                    env=env_variables,
+                    check=False,
+                    stdout=f,
+                    stderr=subprocess.STDOUT,
+                )
+            finally:
+                thread.stop()
+                thread.join()
+        with ci_group(f"Result of {title}", message_type=message_type_from_return_code(result.returncode)):
+            with open(f.name) as f:
+                shutil.copyfileobj(f, sys.stdout)
+    finally:
+        os.unlink(f.name)
+    return result
+
+
 @main.command(
     name='tests',
     help="Run the specified unit test targets. Multiple targets may be specified separated by spaces.",
@@ -122,10 +224,19 @@ def docker_compose_tests(
 )
 @option_dry_run
 @option_verbose
+@option_python
+@option_backend
+@option_postgres_version
+@option_mysql_version
+@option_mssql_version
 @option_integration
+@click.option(
+    '--limit-progress-output',
+    help="Limit progress to percentage only and just show the summary when tests complete.",
+    is_flag=True,
+)
 @click.argument('extra_pytest_args', nargs=-1, type=click.UNPROCESSED)
 @click.option(
-    "-tt",
     "--test-type",
     help="Type of test to run.",
     default="All",
@@ -135,6 +246,12 @@ def docker_compose_tests(
 def tests(
     dry_run: bool,
     verbose: bool,
+    python: str,
+    backend: str,
+    postgres_version: str,
+    mysql_version: str,
+    mssql_version: str,
+    limit_progress_output: bool,
     integration: Tuple,
     extra_pytest_args: Tuple,
     test_type: str,
@@ -149,11 +266,39 @@ def tests(
         os.environ["LIST_OF_INTEGRATION_TESTS_TO_RUN"] = ' '.join(list(integration))
     if db_reset:
         os.environ["DB_RESET"] = "true"
-
-    exec_shell_params = ShellParams(verbose=verbose, dry_run=dry_run)
+    exec_shell_params = ShellParams(
+        verbose=verbose,
+        dry_run=dry_run,
+        python=python,
+        backend=backend,
+        postgres_version=postgres_version,
+        mysql_version=mysql_version,
+        mssql_version=mssql_version,
+    )
     env_variables = get_env_variables_for_docker_commands(exec_shell_params)
     perform_environment_checks(verbose=verbose)
     cmd = ['docker-compose', 'run', '--service-ports', '--rm', 'airflow']
     cmd.extend(list(extra_pytest_args))
-    result = run_command(cmd, verbose=verbose, dry_run=dry_run, env=env_variables, check=False)
+    version = (
+        mssql_version
+        if backend == "mssql"
+        else mysql_version
+        if backend == "mysql"
+        else postgres_version
+        if backend == "postgres"
+        else "none"
+    )
+    if limit_progress_output:
+        result = run_with_progress(
+            cmd=cmd,
+            env_variables=env_variables,
+            test_type=test_type,
+            python=python,
+            backend=backend,
+            version=version,
+            verbose=verbose,
+            dry_run=dry_run,
+        )
+    else:
+        result = run_command(cmd, verbose=verbose, dry_run=dry_run, env=env_variables, check=False)
     sys.exit(result.returncode)

dev/breeze/src/airflow_breeze/utils/ci_group.py

@@ -18,11 +18,11 @@
 import os
 from contextlib import contextmanager
 
-from airflow_breeze.utils.console import get_console
+from airflow_breeze.utils.console import MessageType, get_console
 
 
 @contextmanager
-def ci_group(title: str, enabled: bool = True):
+def ci_group(title: str, enabled: bool = True, message_type: MessageType = MessageType.INFO):
     """
     If used in GitHub Action, creates an expandable group in the GitHub Action log.
     Otherwise, display simple text groups.
@@ -34,9 +34,9 @@ def ci_group(title: str, enabled: bool = True):
         yield
         return
     if os.environ.get('GITHUB_ACTIONS', 'false') != "true":
-        get_console().print(f"[info]{title}[/]")
+        get_console().print(f"[{message_type.value}]{title}[/]")
         yield
         return
-    get_console().print(f"::group::<CLICK_TO_EXPAND>: [info]{title}[/]")
+    get_console().print(f"::group::<CLICK_TO_EXPAND>: [{message_type.value}]{title}[/]")
     yield
     get_console().print("::endgroup::")

dev/breeze/src/airflow_breeze/utils/console.py

@@ -19,6 +19,7 @@
 to be only run in CI or real development terminal - in both cases we want to have colors on.
 """
 import os
+from enum import Enum
 from functools import lru_cache
 
 from rich.console import Console
@@ -56,6 +57,19 @@ def get_theme() -> Theme:
     )
 
 
+class MessageType(Enum):
+    SUCCESS = "success"
+    INFO = "info"
+    WARNING = "warning"
+    ERROR = "error"
+
+
+def message_type_from_return_code(return_code: int) -> MessageType:
+    if return_code == 0:
+        return MessageType.SUCCESS
+    return MessageType.ERROR
+
+
 @lru_cache(maxsize=None)
 def get_console() -> Console:
     return Console(

images/breeze/output-commands-hash.txt

@@ -1 +1 @@
-7f2019004f86eeab48332eb0ea11114d
+2942c0bca323521e3e9af5922d527201