From f578e8d25bb624785e4b1511e0daaea2ef6af79f Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 7 Apr 2025 04:34:12 +0530
Subject: [PATCH] Cleanup various test related things

---
 kitty_tests/__init__.py | 25 +++++++++++++++++--------
 kitty_tests/ssh.py      |  2 --
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/kitty_tests/__init__.py b/kitty_tests/__init__.py
index 21fedbe61..cda1a38df 100644
--- a/kitty_tests/__init__.py
+++ b/kitty_tests/__init__.py
@@ -215,18 +215,22 @@ def filled_history_buf(ynum=5, xnum=5, cursor=Cursor()):
     return ans
 
 
-def retry_on_failure(max_attempts=2, sleep_duration=2):
+is_ci = os.environ.get('CI') == 'true'
+
+
+def retry_on_failure(max_attempts=4 if is_ci else 2, sleep_duration=30 if is_ci else 2):
     def decorator(func):
         @wraps(func)
         def wrapper(*args, **kwargs):
             for attempt in range(max_attempts):
                 try:
                     return func(*args, **kwargs)
-                except Exception as e:
+                except Exception:
                     if attempt < max_attempts - 1: # Don't sleep on the last attempt
                         time.sleep(sleep_duration)
+                        print(f'{func.__name__} failed, retrying in {sleep_duration} seconds', file=sys.stderr)
                     else:
-                        raise e # Re-raise the last exception
+                        raise # Re-raise the last exception
         return wrapper
     return decorator
 
@@ -235,7 +239,7 @@ class BaseTest(TestCase):
 
     ae = TestCase.assertEqual
     maxDiff = 2048
-    is_ci = os.environ.get('CI') == 'true'
+    is_ci = is_ci
 
     def rmtree_ignoring_errors(self, tdir):
         try:
@@ -401,13 +405,13 @@ class PTY:
                 self.process_input_from_child(timeout=end_time - time.monotonic())
             except OSError as e:
                 if not q():
-                    raise Exception(f'Failed to read from pty with error: {e}. Screen contents:  \n {repr(self.screen_contents())}') from e
+                    raise Exception(f'Failed to read from pty with error: {e}. {self.screen_contents_for_error()}') from e
                 return
         if not q():
             msg = 'The condition was not met'
             if timeout_msg is not None:
                 msg = timeout_msg()
-            raise TimeoutError(f'Timed out after {timeout} seconds: {msg}. Screen contents: \n {repr(self.screen_contents())}')
+            raise TimeoutError(f'Timed out after {timeout} seconds: {msg}. {self.screen_contents_for_error()}')
 
     def wait_till_child_exits(self, timeout=30 if BaseTest.is_ci else 10, require_exit_code=None):
         end_time = time.monotonic() + timeout
@@ -419,11 +423,11 @@ class PTY:
                 if require_exit_code is not None and ec != require_exit_code:
                     raise AssertionError(
                         f'Child exited with exit status: {status} code: {ec} != {require_exit_code}.'
-                        f' Screen contents:\n{self.screen_contents()}')
+                        f' {self.screen_contents_for_error()}')
                 return status
             with suppress(OSError):
                 self.process_input_from_child(timeout=0.02)
-        raise AssertionError(f'Child did not exit in {timeout} seconds. Screen contents:\n{self.screen_contents()}')
+        raise AssertionError(f'Child did not exit in {timeout} seconds. {self.screen_contents_for_error()}')
 
     def set_window_size(self, rows=25, columns=80, send_signal=True):
         if hasattr(self, 'screen'):
@@ -434,6 +438,11 @@ class PTY:
             s = struct.pack('HHHH', rows, columns, x_pixels, y_pixels)
             fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, s)
 
+    def screen_contents_for_error(self):
+        from kitty.window import as_text
+        ans = as_text(self.screen, add_history=True, as_ansi=False)
+        return f'Screen contents as repr:\n{ans!r}\nScreen contents:\n{ans.rstrip()}'
+
     def screen_contents(self):
         lines = []
         for i in range(self.screen.lines):
diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py
index 150e303a7..f03c63d03 100644
--- a/kitty_tests/ssh.py
+++ b/kitty_tests/ssh.py
@@ -255,8 +255,6 @@ env COLORTERM
 
             def check_untar_or_fail():
                 q = pty.screen_contents()
-                if 'bzip2' in q:
-                    raise ValueError('Untarring failed with screen contents:\n' + q)
                 return 'UNTAR_DONE' in q
             pty.wait_till(check_untar_or_fail, timeout=60)
             self.assertTrue(os.path.exists(os.path.join(home_dir, '.terminfo/kitty.terminfo')))