From 2c9a4dd5ad643562b83c5c8af72c13536f280e75 Mon Sep 17 00:00:00 2001
From: Sorin Sbarnea <sorin.sbarnea@gmail.com>
Date: Mon, 28 Oct 2024 13:01:36 +0000
Subject: [PATCH] Improve test failure output (#450)

---
 .pre-commit-config.yaml       |  2 +-
 pyproject.toml                |  1 +
 tests/conftest.py             | 25 ++++++++++++++++++++++++-
 tests/integration/conftest.py | 18 ++++++++++++------
 4 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 3d30122d..4f877e42 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -68,7 +68,7 @@ repos:
       - id: tox-ini-fmt
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.7.0
+    rev: v0.7.1
     hooks:
       - id: ruff
         args:
diff --git a/pyproject.toml b/pyproject.toml
index a9a64707..a7457a01 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -335,6 +335,7 @@ builtins = ["__"]
 cache-dir = "./.cache/.ruff"
 fix = true
 line-length = 100
+required-version = ">=0.7.1"
 target-version = "py310"
 
 [tool.ruff.lint]
diff --git a/tests/conftest.py b/tests/conftest.py
index 59634cc3..93fa0b6d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -22,6 +22,7 @@
 import errno
 import os
 import pty
+import re
 import select
 import shutil
 import subprocess
@@ -371,6 +372,24 @@ def _stop_container() -> None:
     subprocess.run(cmd, check=True, capture_output=True)  # noqa: S603
 
 
+def strip_ansi_escape(data: str | bytes) -> str:
+    """Remove all ANSI escapes from string or bytes.
+
+    If bytes is passed instead of string, it will be converted to string
+    using UTF-8.
+
+    Args:
+        data: Input string or bytes
+
+    Returns:
+        String with removed ANSI sequences
+    """
+    if isinstance(data, bytes):  # pragma: no branch
+        data = data.decode("utf-8")
+
+    return re.sub(r"\x1b[^m]*m", "", data)
+
+
 def _exec_container(command: str) -> subprocess.CompletedProcess[str]:
     """Run the container.
 
@@ -384,13 +403,17 @@ def _exec_container(command: str) -> subprocess.CompletedProcess[str]:
         f"{INFRASTRUCTURE.container_engine} exec -t"
         f" {INFRASTRUCTURE.container_name} bash -c '{command}'"
     )
-    return subprocess.run(
+    result = subprocess.run(
         cmd,
         check=False,
         capture_output=True,
         text=True,
         shell=True,
     )
+    # we need to strip the ansi escape codes from the output
+    result.stdout = strip_ansi_escape(result.stdout)
+    result.stderr = strip_ansi_escape(result.stderr)
+    return result
 
 
 @pytest.fixture
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index e60801fb..d8872343 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import os
 import time
 
 from typing import TYPE_CHECKING
@@ -51,7 +52,8 @@ def _shell_in_container(self) -> None:
             f"{self.infrastructure.container_engine} exec -it"
             f" {self.infrastructure.container_name} /bin/zsh"
         )
-        self.send_and_wait(cmd, "workdir", 3)
+        # 3s was not enough as we seen the arm64 runner timing out often
+        self.send_and_wait(cmd, "workdir", 6)
 
     def send_and_wait(self, cmd: str, wait_for: str, timeout: float = 3.0) -> list[str]:
         """Send a command and wait for a response.
@@ -69,16 +71,20 @@ def send_and_wait(self, cmd: str, wait_for: str, timeout: float = 3.0) -> list[s
             err = "No active pane found."
             pytest.fail(err)
         self.pane.send_keys(cmd)
-        timeout = time.time() + timeout
+        start_time = time.time()
         while True:
             stdout = self.pane.capture_pane()
+            stdout_list = stdout if isinstance(stdout, list) else [stdout]
             if not wait_for:
-                return stdout if isinstance(stdout, list) else [stdout]
+                return stdout_list
             if any(wait_for in line for line in stdout):
-                return stdout if isinstance(stdout, list) else [stdout]
-            if time.time() > timeout:
+                return stdout_list
+            if time.time() > timeout + start_time:
                 break
-        error = f"Timeout waiting for {wait_for} in {stdout}"
+        error = (
+            f"Timeout waiting for {timeout} seconds to find string '{wait_for}' in:\n"
+            f" {os.linesep.join(stdout_list)}"
+        )
         pytest.fail(error)
 
     def exit(self) -> None: