From e06bd6837905451ff115dcf5db79b03b028e18ef Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sun, 27 Feb 2022 14:55:49 +0530
Subject: [PATCH] Split up ssh tests

---
 kitty_tests/main.py |  9 ++++++---
 kitty_tests/ssh.py  | 33 +++++++++++++++++----------------
 2 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/kitty_tests/main.py b/kitty_tests/main.py
index 64ade8fe2..91fc63d3a 100644
--- a/kitty_tests/main.py
+++ b/kitty_tests/main.py
@@ -5,10 +5,8 @@ import importlib
 import os
 import sys
 import unittest
-from typing import Callable, Generator, NoReturn, Sequence, Set
-
-
 from importlib.resources import contents
+from typing import Callable, Generator, NoReturn, Sequence, Set
 
 
 def itertests(suite: unittest.TestSuite) -> Generator[unittest.TestCase, None, None]:
@@ -88,10 +86,15 @@ def run_tests() -> None:
         'name', nargs='*', default=[],
         help='The name of the test to run, for e.g. linebuf corresponds to test_linebuf. Can be specified multiple times')
     parser.add_argument('--verbosity', default=4, type=int, help='Test verbosity')
+    parser.add_argument('--module', default='', help='Name of a test module to restrict to. For example: ssh')
     args = parser.parse_args()
     if args.name and args.name[0] in ('type-check', 'type_check', 'mypy'):
         type_check()
     tests = find_all_tests()
+    if args.module:
+        tests = filter_tests_by_module(tests, args.module)
+        if not tests._tests:
+            raise SystemExit('No test module named %s found' % args.module)
     if args.name:
         tests = filter_tests_by_name(tests, *args.name)
         if not tests._tests:
diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py
index afcfd8df8..4b97f8489 100644
--- a/kitty_tests/ssh.py
+++ b/kitty_tests/ssh.py
@@ -20,7 +20,7 @@ from . import BaseTest
 from .shell_integration import bash_ok, basic_shell_env
 
 
-class SSHTest(BaseTest):
+class SSHKitten(BaseTest):
 
     def test_basic_pty_operations(self):
         pty = self.create_pty('echo hello')
@@ -72,19 +72,20 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
     def all_possible_sh(self):
         return tuple(sh for sh in ('dash', 'zsh', 'bash', 'posh', 'sh') if shutil.which(sh))
 
-    def test_ssh_bootstrap_script(self):
-        # test setting env vars
-        with tempfile.TemporaryDirectory() as tdir:
-            pty = self.check_bootstrap(
-                'dash', tdir, extra_exec='env; exit 0', SHELL_INTEGRATION_VALUE='',
-                ssh_opts={'env': {
-                    'TSET': 'set-works',
-                    'COLORTERM': DELETE_ENV_VAR,
-                }}
-            )
-            pty.wait_till(lambda: 'TSET=set-works' in pty.screen_contents())
-            self.assertNotIn('COLORTERM', pty.screen_contents())
-        # test handling of data in tty before tarfile is sent
+    def test_ssh_env_vars(self):
+        for sh in self.all_possible_sh:
+            with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir:
+                pty = self.check_bootstrap(
+                    sh, tdir, extra_exec='env; exit 0', SHELL_INTEGRATION_VALUE='',
+                    ssh_opts={'env': {
+                        'TSET': 'set-works',
+                        'COLORTERM': DELETE_ENV_VAR,
+                    }}
+                )
+                pty.wait_till(lambda: 'TSET=set-works' in pty.screen_contents())
+                self.assertNotIn('COLORTERM', pty.screen_contents())
+
+    def test_ssh_leading_data(self):
         for sh in self.all_possible_sh:
             with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir:
                 pty = self.check_bootstrap(
@@ -92,7 +93,7 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
                     SHELL_INTEGRATION_VALUE='', pre_data='before_tarfile')
                 self.ae(pty.screen_contents(), 'UNTAR_DONE\nld:before_tarfile')
 
-        # test various detection methods for login_shell
+    def test_ssh_login_shell_detection(self):
         methods = []
         if shutil.which('python') or shutil.which('python3') or shutil.which('python2'):
             methods.append('using_python')
@@ -112,7 +113,7 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
                     pty = self.check_bootstrap(sh, tdir, extra_exec=f'{m}; echo "$login_shell"; exit 0', SHELL_INTEGRATION_VALUE='')
                     self.assertIn(expected_login_shell, pty.screen_contents())
 
-        # check that shell integration works
+    def test_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):