From 705421bc5d326952325c4bcde0fd0584de628b4d Mon Sep 17 00:00:00 2001 From: Ethan Date: Fri, 22 May 2026 20:10:25 +1000 Subject: [PATCH] test: speed up agent container websocket close test (#25559) `TestWatchAgentContainers/CoderdWebSocketCanHandleClientClosing` spent about 15 seconds waiting for the real websocket heartbeat ticker to detect that the client closed. Add a clock-aware `HeartbeatClose` wrapper and pass `api.Clock` through the containers watch handler so the test can drive the heartbeat deterministically with `quartz.Mock`. The test still verifies the same client-close teardown path, but it advances the heartbeat tick instead of waiting for wall-clock time. Refs #25557 Discovered as part of the work on CODAGT-381. --- coderd/httpapi/websocket.go | 6 ++++++ coderd/workspaceagents.go | 2 +- coderd/workspaceagents_internal_test.go | 13 ++++++++++++- 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/coderd/httpapi/websocket.go b/coderd/httpapi/websocket.go index c483cf1834..767007aa8e 100644 --- a/coderd/httpapi/websocket.go +++ b/coderd/httpapi/websocket.go @@ -21,6 +21,12 @@ func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn * heartbeatCloseWith(ctx, logger, exit, conn, quartz.NewReal(), HeartbeatInterval) } +// HeartbeatCloseWithClock is like HeartbeatClose, but uses the provided +// clock so tests can drive heartbeat ticks deterministically. +func HeartbeatCloseWithClock(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn, clk quartz.Clock) { + heartbeatCloseWith(ctx, logger, exit, conn, clk, HeartbeatInterval) +} + func heartbeatCloseWith(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn, clk quartz.Clock, interval time.Duration) { ticker := clk.NewTicker(interval, "HeartbeatClose") defer ticker.Stop() diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index 9f830d4f40..b3074740b1 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -871,7 +871,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageText) defer wsNetConn.Close() - go httpapi.HeartbeatClose(ctx, logger, cancel, conn) + go httpapi.HeartbeatCloseWithClock(ctx, logger, cancel, conn, api.Clock) encoder := json.NewEncoder(wsNetConn) diff --git a/coderd/workspaceagents_internal_test.go b/coderd/workspaceagents_internal_test.go index ccbfff4330..a3ff57f025 100644 --- a/coderd/workspaceagents_internal_test.go +++ b/coderd/workspaceagents_internal_test.go @@ -26,6 +26,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac/policy" @@ -37,6 +38,7 @@ import ( "github.com/coder/coder/v2/tailnet" "github.com/coder/coder/v2/tailnet/tailnettest" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" "github.com/coder/websocket" ) @@ -738,8 +740,9 @@ func TestWatchAgentContainers(t *testing.T) { // response to this issue: https://github.com/coder/coder/issues/19449 var ( - ctx = testutil.Context(t, testutil.WaitLong) + ctx = testutil.Context(t, testutil.WaitShort) logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug).Named("coderd") + mClock = quartz.NewMock(t) mCtrl = gomock.NewController(t) mDB = dbmock.NewMockStore(mCtrl) @@ -766,12 +769,16 @@ func TestWatchAgentContainers(t *testing.T) { AgentInactiveDisconnectTimeout: testutil.WaitShort, Database: mDB, Logger: logger, + Clock: mClock, DeploymentValues: &codersdk.DeploymentValues{}, TailnetCoordinator: tailnettest.NewFakeCoordinator(), }, } ) + trap := mClock.Trap().NewTicker("HeartbeatClose") + defer trap.Close() + var tailnetCoordinator tailnet.Coordinator = mCoordinator api.TailnetCoordinator.Store(&tailnetCoordinator) api.agentProvider = fAgentProvider @@ -817,6 +824,8 @@ func TestWatchAgentContainers(t *testing.T) { defer resp.Body.Close() } + trap.MustWait(ctx).MustRelease(ctx) + // And: Create a streaming decoder decoder := wsjson.NewDecoder[codersdk.WorkspaceAgentListContainersResponse](conn, websocket.MessageText, logger) defer decoder.Close() @@ -836,6 +845,7 @@ func TestWatchAgentContainers(t *testing.T) { // When: We close the WebSocket conn.Close(websocket.StatusNormalClosure, "test closing connection") + mClock.Advance(httpapi.HeartbeatInterval).MustWait(ctx) // Then: We expect `containersCh` to be closed. select { @@ -883,6 +893,7 @@ func TestWatchAgentContainers(t *testing.T) { AgentInactiveDisconnectTimeout: testutil.WaitShort, Database: mDB, Logger: logger, + Clock: quartz.NewReal(), DeploymentValues: &codersdk.DeploymentValues{}, TailnetCoordinator: tailnettest.NewFakeCoordinator(), },