Skip to content

Commit 0035049

Browse files
spikecurtisspikecurtis
committed
feat: add cleanup to task-status load test runner
1 parent 83b6690 commit 0035049

File tree

4 files changed

+126
-1
lines changed

4 files changed

+126
-1
lines changed

‎cli/exp_scaletest_taskstatus.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,17 @@ After all runners connect, it waits for the baseline duration before triggering
206206
}
207207
}
208208

209+
cleanupCtx, cleanupCancel := cleanupStrategy.toContext(ctx)
210+
defer cleanupCancel()
211+
err = th.Cleanup(cleanupCtx)
212+
if err != nil {
213+
return xerrors.Errorf("cleanup tests: %w", err)
214+
}
215+
216+
if res.TotalFail > 0 {
217+
return xerrors.New("load test failed, see above for more details")
218+
}
219+
209220
return nil
210221
},
211222
}

‎scaletest/taskstatus/client.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ type client interface {
3434
// watchWorkspace watches for updates to a workspace.
3535
watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (<-chan codersdk.Workspace, error)
3636

37+
// deleteWorkspace deletes the workspace by creating a build with delete transition.
38+
deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error
39+
3740
// initialize sets up the client with the provided logger, which is only available after Run() is called.
3841
initialize(logger slog.Logger)
3942
}
@@ -101,6 +104,18 @@ func (c *sdkClient) watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (
101104
return c.coderClient.WatchWorkspace(ctx, workspaceID)
102105
}
103106

107+
func (c *sdkClient) deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
108+
// Create a build with delete transition to delete the workspace
109+
_, err := c.coderClient.CreateWorkspaceBuild(ctx, workspaceID, codersdk.CreateWorkspaceBuildRequest{
110+
Transition: codersdk.WorkspaceTransitionDelete,
111+
Reason: codersdk.CreateWorkspaceBuildReasonCLI,
112+
})
113+
if err != nil {
114+
return xerrors.Errorf("create delete build: %w", err)
115+
}
116+
return nil
117+
}
118+
104119
func (c *sdkClient) initialize(logger slog.Logger) {
105120
// Configure the coder client logging
106121
c.coderClient.SetLogger(logger)

‎scaletest/taskstatus/run.go‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,10 @@ type Runner struct {
4141
clock quartz.Clock
4242
}
4343

44-
var _ harness.Runnable = &Runner{}
44+
var (
45+
_ harness.Runnable = &Runner{}
46+
_ harness.Cleanable = &Runner{}
47+
)
4548

4649
// NewRunner creates a new Runner with the provided codersdk.Client and configuration.
4750
func NewRunner(coderClient *codersdk.Client, cfg Config) *Runner {
@@ -111,6 +114,30 @@ func (r *Runner) Run(ctx context.Context, name string, logs io.Writer) error {
111114
return nil
112115
}
113116

117+
// Cleanup deletes the external workspace created by this runner.
118+
func (r *Runner) Cleanup(ctx context.Context, id string, logs io.Writer) error {
119+
if r.workspaceID == uuid.Nil {
120+
// No workspace was created, nothing to cleanup
121+
return nil
122+
}
123+
124+
logs = loadtestutil.NewSyncWriter(logs)
125+
logger := slog.Make(sloghuman.Sink(logs)).Leveled(slog.LevelDebug).Named(id)
126+
127+
logger.Info(ctx, "deleting external workspace", slog.F("workspace_id", r.workspaceID))
128+
129+
err := r.client.deleteWorkspace(ctx, r.workspaceID)
130+
if err != nil {
131+
logger.Error(ctx, "failed to delete external workspace",
132+
slog.F("workspace_id", r.workspaceID),
133+
slog.Error(err))
134+
return xerrors.Errorf("delete external workspace: %w", err)
135+
}
136+
137+
logger.Info(ctx, "successfully deleted external workspace", slog.F("workspace_id", r.workspaceID))
138+
return nil
139+
}
140+
114141
func (r *Runner) watchWorkspaceUpdates(ctx context.Context) error {
115142
shouldMarkConnectedDone := true
116143
defer func() {

‎scaletest/taskstatus/run_internal_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"golang.org/x/xerrors"
1515

1616
"cdr.dev/slog"
17+
"cdr.dev/slog/sloggers/sloghuman"
1718
"github.com/coder/quartz"
1819

1920
"github.com/coder/coder/v2/codersdk"
@@ -57,6 +58,12 @@ func (m *fakeClient) createExternalWorkspace(ctx context.Context, req codersdk.C
5758
}, nil
5859
}
5960

61+
func (m *fakeClient) deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
62+
m.logger.Debug(ctx, "called fake DeleteWorkspace", slog.F("workspace_id", workspaceID.String()))
63+
// Simulate successful deletion in tests
64+
return nil
65+
}
66+
6067
// fakeAppStatusPatcher implements the appStatusPatcher interface for testing
6168
type fakeAppStatusPatcher struct {
6269
t *testing.T
@@ -480,3 +487,68 @@ func TestParseStatusMessage(t *testing.T) {
480487
})
481488
}
482489
}
490+
491+
func TestRunner_Cleanup(t *testing.T) {
492+
t.Parallel()
493+
494+
ctx := testutil.Context(t, testutil.WaitMedium)
495+
496+
fakeClient := &fakeClientWithCleanupTracking{
497+
fakeClient: newFakeClient(t),
498+
deleteWorkspaceCalls: make([]uuid.UUID, 0),
499+
}
500+
fakeClient.initialize(slog.Make(sloghuman.Sink(testutil.NewTestLogWriter(t))).Leveled(slog.LevelDebug))
501+
502+
cfg := Config{
503+
AppSlug: "test-app",
504+
TemplateID: uuid.UUID{5, 6, 7, 8},
505+
WorkspaceName: "test-workspace",
506+
MetricLabelValues: []string{"test"},
507+
Metrics: NewMetrics(prometheus.NewRegistry(), "test"),
508+
ReportStatusPeriod: 100 * time.Millisecond,
509+
ReportStatusDuration: 200 * time.Millisecond,
510+
StartReporting: make(chan struct{}),
511+
ConnectedWaitGroup: &sync.WaitGroup{},
512+
}
513+
514+
runner := &Runner{
515+
client: fakeClient,
516+
patcher: newFakeAppStatusPatcher(t),
517+
cfg: cfg,
518+
clock: quartz.NewMock(t),
519+
}
520+
521+
logWriter := testutil.NewTestLogWriter(t)
522+
523+
// Case 1: No workspace created - Cleanup should do nothing
524+
err := runner.Cleanup(ctx, "test-runner", logWriter)
525+
require.NoError(t, err)
526+
require.Len(t, fakeClient.deleteWorkspaceCalls, 0, "deleteWorkspace should not be called when no workspace was created")
527+
528+
// Case 2: Workspace created - Cleanup should delete it
529+
runner.workspaceID = uuid.UUID{1, 2, 3, 4}
530+
err = runner.Cleanup(ctx, "test-runner", logWriter)
531+
require.NoError(t, err)
532+
require.Len(t, fakeClient.deleteWorkspaceCalls, 1, "deleteWorkspace should be called once")
533+
require.Equal(t, runner.workspaceID, fakeClient.deleteWorkspaceCalls[0], "deleteWorkspace should be called with correct workspace ID")
534+
535+
// Case 3: Cleanup with error
536+
fakeClient.deleteError = xerrors.New("delete failed")
537+
runner.workspaceID = uuid.UUID{5, 6, 7, 8}
538+
err = runner.Cleanup(ctx, "test-runner", logWriter)
539+
require.Error(t, err)
540+
require.Contains(t, err.Error(), "delete external workspace")
541+
}
542+
543+
// fakeClientWithCleanupTracking extends fakeClient to track deleteWorkspace calls
544+
type fakeClientWithCleanupTracking struct {
545+
*fakeClient
546+
deleteWorkspaceCalls []uuid.UUID
547+
deleteError error
548+
}
549+
550+
func (c *fakeClientWithCleanupTracking) deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
551+
c.deleteWorkspaceCalls = append(c.deleteWorkspaceCalls, workspaceID)
552+
c.logger.Debug(ctx, "called fake DeleteWorkspace with tracking", slog.F("workspace_id", workspaceID.String()))
553+
return c.deleteError
554+
}

0 commit comments

Comments
 (0)