acc: converge ssh/connect-serverless-gpu tunnel test onto a local web…#5920
Merged
Conversation
…socket route Follow-up to #5878. That PR covered the SSH tunnel in two styles: an acceptance test asserting the submitted bootstrap job, and a Go test (TestClientServerRealSSHD) driving the tunnel transport. This folds the tunnel coverage into the acceptance test so a single local run exercises both layers. Two enabling changes: - Product fix (experimental/ssh/internal/client/websockets.go): derive the websocket scheme from the workspace host (http -> ws, else wss) instead of hardcoding wss, so the client can dial the plaintext local test server. Adds a unit test for the URL builder. - Test server (libs/testserver): add a raw, connection-hijacking handler registration (Router.HandleRaw) plus a driver-proxy /ssh websocket route that echoes binary frames, standing in for the SSH tunnel server a real cluster runs. In-process, so it also works on Windows. acceptance/ssh/connect-serverless-gpu now runs a single `ssh connect --proxy` locally: it submits the same serverless-GPU bootstrap job and pipes stdin/stdout over the tunnel, asserting both the job payload and the echoed bytes. Proxy mode skips the ssh subprocess, so it doesn't hang looping a handshake against the echo.
Collaborator
Integration test reportCommit: 04c8e74
12 interesting tests: 5 flaky, 4 SKIP, 3 KNOWN
Top 12 slowest tests (at least 2 minutes):
|
…s locally Deepens the convergence started earlier on this branch. Instead of echoing frames, the test server now upgrades the driver-proxy /ssh websocket to a real `sshd -i` and bridges binary frames to its stdio, so a single local `ssh connect ... -- "echo ..."` completes a genuine SSH handshake, public-key auth, and remote command exec with no Databricks compute. - libs/testserver: replace the echo handler with (*Server).sshTunnelHandler, which reads the CLI's uploaded client-public-key secret (the ws dial carries the bearer token, so the token-scoped workspace is the one that stored it), writes an ephemeral host key + authorized_keys + a minimal sshd_config, spawns sshd, and pumps ws frames <-> sshd stdin/stdout. /metadata now returns the current OS user so `ssh -l <user>` matches the account sshd runs as. - acceptance/ssh/connect-serverless-gpu: run the full handshake locally, still assert the bootstrap job payload, and self-skip via SKIP_TEST when sshd is absent (the general test job doesn't provision it; task test-exp-ssh does). Gated to linux, matching the removed Go test's rationale. - Delete TestClientServerRealSSHD: the acceptance test now covers the real handshake end to end, so the Go stand-in is redundant. The cat-echo transport tests remain.
CI caught two issues in the previous commit: - lint: use sync.OnceFunc and wg.Go instead of sync.Once/wg.Add+Done in the test server's ssh handler. - GOOS keys absent from the map default to enabled, so `linux = true` alone did not disable macOS/Windows. They ran the handshake: macOS found its system sshd and hung, Windows found sshd.exe and failed fast. Disable darwin and windows explicitly to keep the test Linux-only.
anton-107
approved these changes
Jul 14, 2026
anton-107
left a comment
anton-107
left a comment
Contributor
There was a problem hiding this comment.
Approve — nicely done. The convergence works: the CLI subprocess drives the real client proxy (RunClientProxy) while the testserver upgrades /ssh to a real sshd -i, so one local run genuinely exercises handshake + pubkey auth + remote exec with no compute. Confirmed it actually runs (not skipped): TestAccept/ssh/connect-serverless-gpu passes in ~0.5s on the Linux test-exp-ssh job. The buildProxyWebsocketURL http→ws / else→wss refactor is correct for all real inputs (Config.Host is always scheme-normalized by the SDK before use).
One non-blocking note on test strategy: the deleted TestClientServerRealSSHD was the only test that carried a real SSH handshake through the production server proxy (proxy.NewProxyServer/runServerProxy). The new libs/testserver/ssh.go handler is a hand-rolled ws↔sshd pump that never calls NewProxyServer, so that specific coverage is lost. It is narrow — runServerProxy is protocol-agnostic and the surviving cat-echo tests still drive it end-to-end for frame pumping / interleaving / handover — what is uniquely gone is the real-handshake banner emission and natural-sshd-exit teardown ordering. Low regression-escape risk, so not a blocker, but worth a sentence in the description that this is a client-side-only real-handshake test. If cheap, having sshTunnelHandler delegate to proxy.NewProxyServer would reclaim the coverage and drop the duplicated pump/findSSHD logic.
Minor, optional: a one-line t.Logf on the early-return failure paths in sshTunnelHandler (missing secret / pipe / Start error) would make a future breakage self-explanatory instead of surfacing only as an empty-stdout golden diff.
The testserver /ssh handler returned silently on every setup failure (missing secret, pipe/Start errors, etc.), so a broken sshd setup only surfaced as an empty-stdout golden diff. Log each early-return reason so a future breakage is self-explanatory.
The testserver /ssh route now drives a real sshd, so ssh/connection's local run completes a genuine handshake and remote exec instead of producing nothing. Mirror connect-serverless-gpu: skip when sshd is absent, gate to Linux (macOS hangs on system sshd, Windows differs), and record "Connection successful" as the golden stdout.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Changes
Follow-up to #5878, converging the SSH tunnel's two test styles into one acceptance test. Two enabling changes: derive the driver-proxy websocket scheme from the host (
http→ws, elsewss) instead of hardcodingwssso the client can dial the plaintext local test server, and add a testserver/sshwebsocket route (via a new connection-hijackingRouter.HandleRaw) that upgrades the driver-proxy request to a realsshd -iand bridges binary frames to its stdio, standing in for the tunnel server a real cluster runs — no Databricks compute required.The
/sshhandler reads the client public key the CLI uploaded viasecrets/put(the ws dial carries the bearer token, so the token-scoped workspace is the one that stored it), writes an ephemeral host key +authorized_keys+ a minimalsshd_config, and pumps websocket frames to and fromsshd's stdin/stdout./metadatanow returns the current OS user sossh -l <user>matches the accountsshdruns as.acceptance/ssh/connect-serverless-gpunow runs a single localssh connect ... -- "echo 'Connection successful'"that both submits the serverless-GPU bootstrap job and completes a genuine SSH handshake, public-key auth, and remote command exec over thews://tunnel, asserting both the job payload and the command output. This makesTestClientServerRealSSHDredundant, so it's removed; thecat-echo transport tests inclient_server_test.goremain.Note the converged test exercises the real handshake on the client side only: the
/sshroute is a testserver stand-in, not the productionproxy.NewProxyServer/runServerProxy, so the server proxy's transport stays covered by thecat-echo tests rather than by a real handshake. This is a deliberate trade-off (runServerProxyis protocol-agnostic), not an accidental coverage drop.The handshake needs a real OpenSSH server: the test self-skips via
SKIP_TESTwhensshdis absent (the generaltestjob doesn't provision it;task test-exp-sshinstallsopenssh-serveron Linux). It's pinned to Linux (darwin/windowsdisabled), where runningsshd -ias a non-root user with a throwaway config is reliable — matching the removed Go test's rationale.Tests
acceptance/ssh/connect-serverless-gpu(bootstrap job payload + real SSH handshake and remote exec over thews://tunnel)TestBuildProxyWebsocketURL(ws/wssscheme selection)