Add cloud-only acceptance tests for ssh tunnel

Author: ilia-db

acceptance/acceptance_test.go

@@ -163,7 +163,7 @@ func testAccept(t *testing.T, inprocessMode bool, singleTest string) int {
 	}
 
 	// Download terraform and provider and create config.
-	RunCommand(t, []string{"python3", filepath.Join(cwd, "install_terraform.py"), "--targetdir", terraformDir}, ".")
+	RunCommand(t, []string{"python3", filepath.Join(cwd, "install_terraform.py"), "--targetdir", terraformDir}, ".", []string{})
 
 	wheelPath := buildDatabricksBundlesWheel(t, buildDir)
 	if wheelPath != "" {
@@ -234,6 +234,12 @@ func testAccept(t *testing.T, inprocessMode bool, singleTest string) int {
 		}
 	}
 
+	if cloudEnv != "" && UseVersion == "" {
+		// Create linux release artifacts, to be used by the cloud-only ssh tunnel tests
+		releasesDir := CreateReleaseArtifacts(t, cwd, "linux")
+		t.Setenv("CLI_RELEASES_DIR", releasesDir)
+	}
+
 	testDefaultWarehouseId := os.Getenv("TEST_DEFAULT_WAREHOUSE_ID")
 	if testDefaultWarehouseId != "" {
 		repls.Set(testDefaultWarehouseId, "[TEST_DEFAULT_WAREHOUSE_ID]")
@@ -892,10 +898,62 @@ func BuildCLI(t *testing.T, buildDir, coverDir string) string {
 		args = append(args, "-buildvcs=false")
 	}
 
-	RunCommand(t, args, "..")
+	RunCommand(t, args, "..", []string{})
 	return execPath
 }
 
+// CreateReleaseArtifacts builds release artifacts for the given OS using amd64 and arm64 architectures,
+// archives them into zip files, and returns the directory containing the release artifacts.
+func CreateReleaseArtifacts(t *testing.T, cwd, osName string) string {
+	releasesDir := filepath.Join(cwd, "build", "releases")
+	require.NoError(t, os.MkdirAll(releasesDir, os.ModePerm))
+	arches := []string{"amd64", "arm64"}
+	for _, arch := range arches {
+		CreateReleaseArtifact(t, cwd, releasesDir, osName, arch)
+	}
+	return releasesDir
+}
+
+func CreateReleaseArtifact(t *testing.T, cwd, releasesDir, osName, arch string) {
+	tempBuildDir := filepath.Join(releasesDir, "tmp_"+arch)
+	require.NoError(t, os.MkdirAll(tempBuildDir, os.ModePerm))
+	defer os.RemoveAll(tempBuildDir)
+
+	execPath := filepath.Join(tempBuildDir, "databricks")
+	args := []string{"go", "build", "-o", execPath}
+	RunCommand(t, args, "..", []string{"GOOS=" + osName, "GOARCH=" + arch})
+
+	zipName := fmt.Sprintf("databricks_cli_%s_%s.zip", osName, arch)
+	zipPath := filepath.Join(releasesDir, zipName)
+
+	zipFile, err := os.Create(zipPath)
+	require.NoError(t, err)
+	defer zipFile.Close()
+
+	zipWriter := zip.NewWriter(zipFile)
+	defer zipWriter.Close()
+
+	info, err := os.Stat(execPath)
+	require.NoError(t, err)
+
+	header, err := zip.FileInfoHeader(info)
+	require.NoError(t, err)
+	header.Name = "databricks"
+	header.Method = zip.Deflate
+
+	writer, err := zipWriter.CreateHeader(header)
+	require.NoError(t, err)
+
+	binaryFile, err := os.Open(execPath)
+	require.NoError(t, err)
+	defer binaryFile.Close()
+
+	_, err = io.Copy(writer, binaryFile)
+	require.NoError(t, err)
+
+	t.Logf("Created %s %s release: %s", osName, arch, zipPath)
+}
+
 // DownloadCLI downloads a released CLI binary archive for the given version,
 // extracts the executable, and returns its path.
 func DownloadCLI(t *testing.T, buildDir, version string) string {
@@ -1116,10 +1174,11 @@ func getUVDefaultCacheDir(t *testing.T) string {
 	}
 }
 
-func RunCommand(t *testing.T, args []string, dir string) {
+func RunCommand(t *testing.T, args []string, dir string, env []string) {
 	start := time.Now()
 	cmd := exec.Command(args[0], args[1:]...)
 	cmd.Dir = dir
+	cmd.Env = append(os.Environ(), env...)
 	out, err := cmd.CombinedOutput()
 	elapsed := time.Since(start)
 	t.Logf("%s took %s", args, elapsed)
@@ -1241,7 +1300,7 @@ func buildDatabricksBundlesWheel(t *testing.T, buildDir string) string {
 	// so we prepare here by keeping only one.
 	_ = prepareWheelBuildDirectory(t, buildDir)
 
-	RunCommand(t, []string{"uv", "build", "--no-cache", "-q", "--wheel", "--out-dir", buildDir}, "../python")
+	RunCommand(t, []string{"uv", "build", "--no-cache", "-q", "--wheel", "--out-dir", buildDir}, "../python", []string{})
 
 	latestWheel := prepareWheelBuildDirectory(t, buildDir)
 	if latestWheel == "" {
@@ -1369,7 +1428,7 @@ func BuildYamlfmt(t *testing.T) {
 	args := []string{
 		"make", "-s", "tools/yamlfmt" + exeSuffix,
 	}
-	RunCommand(t, args, "..")
+	RunCommand(t, args, "..", []string{})
 }
 
 func loadUserReplacements(t *testing.T, repls *testdiff.ReplacementsContext, tmpDir string) {

acceptance/ssh/connection/out.test.toml

@@ -0,0 +1 @@
+Connection successful

acceptance/ssh/connection/output.txt

@@ -0,0 +1,8 @@
+errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >LOG.stdout 2>LOG.stderr
+
+cat LOG.stdout
+
+if ! grep -q "Connection successful" LOG.stdout; then
+  run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$")
+  trace $CLI jobs get-run "$run_id" > LOG.job
+fi

acceptance/ssh/connection/script

@@ -0,0 +1,6 @@
+Local = false
+Cloud = true
+RequiresCluster = true
+
+[EnvMatrix]
+  DATABRICKS_BUNDLE_ENGINE = ["direct-exp"]