shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-02 20:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore(coderd): use WorkspaceAgentWaiter.WithContext in aitasks_test.go (#20360)

Browse Source 
Fixes https://github.com/coder/internal/issues/1067

- Adds `WorkspaceAgentWaiter.WithContext()`
- Updates usage of `WorkspaceAgentWaiter` in `aitasks_test.go` with
context bumped to `testutil.WaitMedium`

Authored by Claude with manual review and updates.
This commit is contained in:
Cian Johnston
2025-10-17 13:45:11 +01:00
committed by GitHub
parent 9ee3a9174c
commit 1a104751e4
2 changed files with 28 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/coderdtest/coderdtest.go
+25 -9
View File
@@ -1128,6 +1128,7 @@ type WorkspaceAgentWaiter struct {
workspaceID uuid.UUID
agentNames []string
resourcesMatcher func([]codersdk.WorkspaceResource) bool
ctx context.Context
}
// NewWorkspaceAgentWaiter returns an object that waits for agents to connect when
@@ -1156,6 +1157,14 @@ func (w WorkspaceAgentWaiter) MatchResources(m func([]codersdk.WorkspaceResource
return w
}
// WithContext instructs the waiter to use the provided context for all operations.
// If not specified, the waiter will create its own context with testutil.WaitLong timeout.
func (w WorkspaceAgentWaiter) WithContext(ctx context.Context) WorkspaceAgentWaiter {
//nolint: revive // returns modified struct
w.ctx = ctx
return w
}
// WaitForAgentFn represents a boolean assertion to be made against each agent
// that a given WorkspaceAgentWaited knows about. Each WaitForAgentFn should apply
// the check to a single agent, but it should be named for plural, because `func (w WorkspaceAgentWaiter) WaitFor`
@@ -1176,6 +1185,8 @@ func AgentsNotReady(agent codersdk.WorkspaceAgent) bool {
return !AgentsReady(agent)
}
// WaitFor waits for the given criteria and fails the test if they are not met before the
// waiter's context is canceled.
func (w WorkspaceAgentWaiter) WaitFor(criteria ...WaitForAgentFn) {
w.t.Helper()
@@ -1184,11 +1195,13 @@ func (w WorkspaceAgentWaiter) WaitFor(criteria ...WaitForAgentFn) {
agentNamesMap[name] = struct{}{}
}
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
ctx := w.ctx
if w.ctx == nil {
ctx = testutil.Context(w.t, testutil.WaitLong)
}
w.t.Logf("waiting for workspace agents (workspace %s)", w.workspaceID)
require.Eventually(w.t, func() bool {
testutil.Eventually(ctx, w.t, func(ctx context.Context) bool {
var err error
workspace, err := w.client.Workspace(ctx, w.workspaceID)
if err != nil {
@@ -1216,10 +1229,11 @@ func (w WorkspaceAgentWaiter) WaitFor(criteria ...WaitForAgentFn) {
}
}
return true
}, testutil.WaitLong, testutil.IntervalMedium)
}, testutil.IntervalMedium)
}
// Wait waits for the agent(s) to connect and fails the test if they do not within testutil.WaitLong
// Wait waits for the agent(s) to connect and fails the test if they do not connect before the
// waiter's context is canceled.
func (w WorkspaceAgentWaiter) Wait() []codersdk.WorkspaceResource {
w.t.Helper()
@@ -1228,12 +1242,14 @@ func (w WorkspaceAgentWaiter) Wait() []codersdk.WorkspaceResource {
agentNamesMap[name] = struct{}{}
}
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
ctx := w.ctx
if w.ctx == nil {
ctx = testutil.Context(w.t, testutil.WaitLong)
}
w.t.Logf("waiting for workspace agents (workspace %s)", w.workspaceID)
var resources []codersdk.WorkspaceResource
require.Eventually(w.t, func() bool {
testutil.Eventually(ctx, w.t, func(ctx context.Context) bool {
var err error
workspace, err := w.client.Workspace(ctx, w.workspaceID)
if err != nil {
@@ -1265,7 +1281,7 @@ func (w WorkspaceAgentWaiter) Wait() []codersdk.WorkspaceResource {
return true
}
return w.resourcesMatcher(resources)
}, testutil.WaitLong, testutil.IntervalMedium)
}, testutil.IntervalMedium)
w.t.Logf("got workspace agents (workspace %s)", w.workspaceID)
return resources
}