Get the ssh bootstrap tests fully working

kitty_tests/__init__.py

@@ -235,9 +235,9 @@ class PTY:
         return bytes_read
 
     def wait_till(self, q, timeout=10):
-        st = time.monotonic()
-        while not q() and time.monotonic() - st < timeout:
-            self.process_input_from_child(timeout=timeout - (time.monotonic() - st))
+        end_time = time.monotonic() + timeout
+        while not q() and time.monotonic() <= end_time:
+            self.process_input_from_child(timeout=max(0, end_time - time.monotonic()))
         if not q():
             raise TimeoutError(f'The condition was not met. Screen contents: \n {repr(self.screen_contents())}')
 

kitty_tests/ssh.py

@@ -9,6 +9,7 @@ import tempfile
 
 from kittens.ssh.main import bootstrap_script, get_connection_data
 from kitty.utils import SSHConnectionData
+from kitty.fast_data_types import CURSOR_BEAM
 
 from . import BaseTest
 from .shell_integration import bash_ok, basic_shell_env
@@ -46,12 +47,24 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
             if q:
                 if sh == 'bash' and not bash_ok():
                     continue
-                with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir:
-                    script = bootstrap_script(EXEC_CMD='echo UNTAR_DONE; exit 0')
-                    env = basic_shell_env(tdir)
-                    pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=tdir, env=env)
-                    self.check_bootstrap(tdir, pty)
+                for login_shell in ('', 'bash', 'zsh', 'fish'):
+                    if (login_shell and not shutil.which(login_shell)) or (login_shell == 'bash' and not bash_ok()):
+                        continue
+                    with self.subTest(sh=sh, login_shell=login_shell), tempfile.TemporaryDirectory() as tdir:
+                        self.check_bootstrap(sh, tdir, login_shell)
 
-    def check_bootstrap(self, home_dir, pty):
-        pty.wait_till(lambda: 'UNTAR_DONE' in pty.screen_contents())
+    def check_bootstrap(self, sh, home_dir, login_shell):
+        ps1 = 'prompt> '
+        script = bootstrap_script(
+            EXEC_CMD=f'echo "UNTAR_DONE"; export PS1="{ps1}"',
+            OVERRIDE_LOGIN_SHELL=login_shell,
+        )
+        env = basic_shell_env(home_dir)
+        # Avoid generating unneeded completion scripts
+        os.makedirs(os.path.join(home_dir, '.local', 'share', 'fish', 'generated_completions'), exist_ok=True)
+        # prevent newuser-install from running
+        open(os.path.join(home_dir, '.zshrc'), 'w').close()
+        self.assertFalse(os.path.exists(os.path.join(home_dir, '.terminfo/kitty.terminfo')))
+        pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=home_dir, env=env)
+        pty.wait_till(lambda: pty.screen.cursor.shape == CURSOR_BEAM)
         self.assertTrue(os.path.exists(os.path.join(home_dir, '.terminfo/kitty.terminfo')))

shell-integration/ssh/bootstrap.sh

@@ -176,12 +176,12 @@ exec_zsh_with_integration() {
     zdotdir="$ZDOTDIR"
     if [ -z "$zdotdir" ]; then 
         zdotdir=~; 
-        unset KITTY_ORIG_ZDOTDIR
+        unset KITTY_ORIG_ZDOTDIR  # ensure this is not propagated
     else
         export KITTY_ORIG_ZDOTDIR="$zdotdir"
     fi
     # dont prevent zsh-new-user from running
-    if [ -e "$zdotdir/.zshrc" || -e "$zdotdir/.zshenv" || -e "$zdotdir/.zprofile" || -e "$zdotdir/.zlogin" ]; then
+    if [ -f "$zdotdir/.zshrc" -o -f "$zdotdir/.zshenv" -o -f "$zdotdir/.zprofile" -o -f "$zdotdir/.zlogin" ]; then
         export ZDOTDIR="$shell_integration_dir/zsh"
         exec "$login_shell" "-l"
     fi
@@ -198,14 +198,14 @@ exec_fish_with_integration() {
 }
 
 exec_with_shell_integration() {
-    case "$login_shell" in
-        *"zsh")
+    case "$shell_name" in
+        "zsh")
             exec_zsh_with_integration
             ;;
-        *"bash")
+        "bash")
             exec_bash_with_integration
             ;;
-        *"fish")
+        "fish")
             exec_fish_with_integration
             ;;
     esac