From e02ba7f389cd6f11e22abb63653906275b71ae2d Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Fri, 24 Feb 2023 08:16:53 +0530
Subject: [PATCH] Port bootstrap script length limit

---
 kitty_tests/ssh.py           |  6 ------
 tools/cmd/ssh/config_test.go |  4 ++--
 tools/cmd/ssh/main_test.go   | 22 ++++++++++++++++++++++
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py
index 7a006ddbb..401b588ac 100644
--- a/kitty_tests/ssh.py
+++ b/kitty_tests/ssh.py
@@ -56,12 +56,6 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         t('ssh --kitten=one -p 12 --kitten two -ix main', identity_file='x', port=12, extra_args=(('--kitten', 'one'), ('--kitten', 'two')))
         self.assertTrue(runtime_dir())
 
-    def test_ssh_bootstrap_sh_cmd_limit(self):
-        # dropbear has a 9000 bytes maximum command length limit
-        sh_script, _, _ = bootstrap_script(SSHOptions({'interpreter': 'sh'}), script_type='sh', remote_args=[], request_id='123-123')
-        rcmd = wrap_bootstrap_script(sh_script, 'sh')
-        self.assertLessEqual(sum(len(x) for x in rcmd), 9000)
-
     @property
     @lru_cache()
     def all_possible_sh(self):
diff --git a/tools/cmd/ssh/config_test.go b/tools/cmd/ssh/config_test.go
index 11e5f2d3d..dff66dbe3 100644
--- a/tools/cmd/ssh/config_test.go
+++ b/tools/cmd/ssh/config_test.go
@@ -27,12 +27,12 @@ func TestSSHConfigParsing(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		actual := c.final_env_instructions(for_python, func(key string) (string, bool) {
+		actual := final_env_instructions(for_python, func(key string) (string, bool) {
 			if key == "LOCAL_ENV" {
 				return "LOCAL_VAL", true
 			}
 			return "", false
-		})
+		}, c.Env...)
 		if expected_env == nil {
 			expected_env = []string{}
 		}
diff --git a/tools/cmd/ssh/main_test.go b/tools/cmd/ssh/main_test.go
index 7d02100f7..c073f0531 100644
--- a/tools/cmd/ssh/main_test.go
+++ b/tools/cmd/ssh/main_test.go
@@ -37,3 +37,25 @@ func TestCloneEnv(t *testing.T) {
 		t.Fatalf("Failed to deserialize env\n%s", diff)
 	}
 }
+
+func basic_connection_data() *connection_data {
+	return &connection_data{
+		script_type: "sh", request_id: "123-123", remote_args: []string{}, host_opts: NewConfig(),
+		username: "testuser", hostname_for_match: "host.test",
+	}
+}
+
+func TestSSHBootstrapScriptLimit(t *testing.T) {
+	cd := basic_connection_data()
+	err := get_remote_command(cd)
+	if err != nil {
+		t.Fatal(err)
+	}
+	total := 0
+	for _, x := range cd.rcmd {
+		total += len(x)
+	}
+	if total > 9000 {
+		t.Fatalf("Bootstrap script too large: %d bytes", total)
+	}
+}