From aca13a619aa8fe3cfc2cec6f84d6bf5ac21b6826 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 26 Mar 2024 19:45:20 +0530
Subject: [PATCH] Retry all the ssh kitten tests on failure once

---
 glfw/wl_window.c        |  2 +-
 kitty_tests/__init__.py | 17 +++++++++++++++++
 kitty_tests/ssh.py      | 18 +++++++++---------
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/glfw/wl_window.c b/glfw/wl_window.c
index 92b2bfd7d..405c90762 100644
--- a/glfw/wl_window.c
+++ b/glfw/wl_window.c
@@ -780,7 +780,7 @@ attach_temp_buffer_during_window_creation(_GLFWwindow *window) {
     else memset(shm_data, 0, size);
     if (!shm_data) {
         close(fd);
-        _glfwInputError(GLFW_PLATFORM_ERROR, "Wayland: failed to mmap anonymouse file");
+        _glfwInputError(GLFW_PLATFORM_ERROR, "Wayland: failed to mmap anonymous file");
         return false;
     }
     struct wl_shm_pool *pool = wl_shm_create_pool(_glfw.wl.shm, fd, size);
diff --git a/kitty_tests/__init__.py b/kitty_tests/__init__.py
index 8716f465e..cbe56ae4b 100644
--- a/kitty_tests/__init__.py
+++ b/kitty_tests/__init__.py
@@ -13,6 +13,7 @@ import sys
 import termios
 import time
 from contextlib import contextmanager
+from functools import wraps
 from pty import CHILD, STDIN_FILENO, STDOUT_FILENO, fork
 from unittest import TestCase
 
@@ -174,6 +175,22 @@ def filled_history_buf(ynum=5, xnum=5, cursor=Cursor()):
     return ans
 
 
+def retry_on_failure(max_attempts=2, sleep_duration=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:
+                    if attempt < max_attempts - 1: # Don't sleep on the last attempt
+                        time.sleep(sleep_duration)
+                    else:
+                        raise e # Re-raise the last exception
+        return wrapper
+    return decorator
+
+
 class BaseTest(TestCase):
 
     ae = TestCase.assertEqual
diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py
index 94f364686..bc4a1d038 100644
--- a/kitty_tests/ssh.py
+++ b/kitty_tests/ssh.py
@@ -8,7 +8,6 @@ import os
 import shutil
 import subprocess
 import tempfile
-import time
 from contextlib import suppress
 from functools import lru_cache
 
@@ -17,7 +16,7 @@ from kitty.constants import is_macos, kitten_exe, runtime_dir
 from kitty.fast_data_types import CURSOR_BEAM, shm_unlink
 from kitty.utils import SSHConnectionData
 
-from . import BaseTest
+from . import BaseTest, retry_on_failure
 from .shell_integration import bash_ok, basic_shell_env
 
 
@@ -29,6 +28,7 @@ def files_in(path):
 
 class SSHKitten(BaseTest):
 
+    @retry_on_failure()
     def test_basic_pty_operations(self):
         pty = self.create_pty('echo hello')
         pty.process_input_from_child()
@@ -41,6 +41,7 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         pty.process_input_from_child()
         self.ae(pty.screen_contents(), '13 77 770 260')
 
+    @retry_on_failure()
     def test_ssh_connection_data(self):
         def t(cmdline, binary='ssh', host='main', port=None, identity_file='', extra_args=()):
             if identity_file:
@@ -63,6 +64,7 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         python = 'python3' if shutil.which('python3') else 'python'
         return tuple(filter(shutil.which, ('dash', 'zsh', 'bash', 'posh', 'sh', python)))
 
+    @retry_on_failure()
     def test_ssh_copy(self):
         simple_data = 'rkjlhfwf9whoaa'
 
@@ -120,6 +122,7 @@ copy --exclude **/w.* --exclude **/r d1
                 })
                 self.ae(len(glob.glob(f'{remote_home}/{tname}/*/xterm-kitty')), 2)
 
+    @retry_on_failure()
     def test_ssh_env_vars(self):
         tset = '$A-$(echo no)-`echo no2` !Q5 "something else"'
         for sh in self.all_possible_sh:
@@ -139,6 +142,7 @@ env COLORTERM
                 pty.wait_till(lambda: '/cwd' in pty.screen_contents())
                 self.assertTrue(pty.is_echo_on())
 
+    @retry_on_failure()
     def test_ssh_bootstrap_with_different_launchers(self):
         for launcher in self.all_possible_sh:
             if 'python' in launcher:
@@ -150,6 +154,7 @@ env COLORTERM
                         with self.subTest(sh=sh, launcher=q), tempfile.TemporaryDirectory() as tdir:
                             self.check_bootstrap(sh, tdir, test_script='env; exit 0', SHELL_INTEGRATION_VALUE='', launcher=q)
 
+    @retry_on_failure()
     def test_ssh_leading_data(self):
         script = 'echo "ld:$leading_data"; exit 0'
         for sh in self.all_possible_sh:
@@ -161,6 +166,7 @@ env COLORTERM
                     SHELL_INTEGRATION_VALUE='', pre_data='before_tarfile')
                 self.ae(pty.screen_contents(), 'UNTAR_DONE\nld:before_tarfile')
 
+    @retry_on_failure()
     def test_ssh_login_shell_detection(self):
         methods = []
         if shutil.which('python') or shutil.which('python3') or shutil.which('python2'):
@@ -188,14 +194,8 @@ env COLORTERM
                     pty = self.check_bootstrap(sh, tdir, test_script=f'{m}; echo "$login_shell"; exit 0', SHELL_INTEGRATION_VALUE='')
                     self.assertIn(expected_login_shell, pty.screen_contents())
 
+    @retry_on_failure()
     def test_ssh_shell_integration(self):
-        try:
-            return self.do_ssh_shell_integration()
-        except Exception:
-            time.sleep(1)
-            self.do_ssh_shell_integration()
-
-    def do_ssh_shell_integration(self):
         ok_login_shell = ''
         for sh in self.all_possible_sh:
             for login_shell in {'fish', 'zsh', 'bash'} & set(self.all_possible_sh):