Skip to content

Commit 7029b02

Browse files
mafredrimafredri
committed
feat(cli): add template filter support to exp scaletest cleanup and traffic
1 parent b83a8ce commit 7029b02

File tree

2 files changed

+131
-32
lines changed

2 files changed

+131
-32
lines changed

‎cli/exp_scaletest.go

+81-32
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ type scaletestStrategyFlags struct {
133133
concurrency int64
134134
timeout time.Duration
135135
timeoutPerJob time.Duration
136+
template string
136137
}
137138

138139
func (s *scaletestStrategyFlags) attach(opts *clibase.OptionSet) {
@@ -394,6 +395,8 @@ func (r *userCleanupRunner) Run(ctx context.Context, _ string, _ io.Writer) erro
394395
}
395396

396397
func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
398+
var template string
399+
397400
cleanupStrategy := &scaletestStrategyFlags{cleanup: true}
398401
client := new(codersdk.Client)
399402

@@ -407,7 +410,7 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
407410
Handler: func(inv *clibase.Invocation) error {
408411
ctx := inv.Context()
409412

410-
_, err := requireAdmin(ctx, client)
413+
me, err := requireAdmin(ctx, client)
411414
if err != nil {
412415
return err
413416
}
@@ -421,8 +424,15 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
421424
},
422425
}
423426

427+
if template != "" {
428+
_, err := parseTemplate(ctx, client, me.OrganizationIDs, template)
429+
if err != nil {
430+
return xerrors.Errorf("parse template: %w", err)
431+
}
432+
}
433+
424434
cliui.Infof(inv.Stdout, "Fetching scaletest workspaces...")
425-
workspaces, err := getScaletestWorkspaces(ctx, client)
435+
workspaces, err := getScaletestWorkspaces(ctx, client, template)
426436
if err != nil {
427437
return err
428438
}
@@ -494,6 +504,15 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
494504
},
495505
}
496506

507+
cmd.Options = clibase.OptionSet{
508+
{
509+
Flag: "template",
510+
Env: "CODER_SCALETEST_CLEANUP_TEMPLATE",
511+
Description: "Name or ID of the template. Only delete workspaces created from the given template.",
512+
Value: clibase.StringOf(&template),
513+
},
514+
}
515+
497516
cleanupStrategy.attach(&cmd.Options)
498517
return cmd
499518
}
@@ -564,34 +583,12 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
564583
return xerrors.Errorf("could not parse --output flags")
565584
}
566585

567-
var tpl codersdk.Template
568586
if template == "" {
569587
return xerrors.Errorf("--template is required")
570588
}
571-
if id, err := uuid.Parse(template); err == nil && id != uuid.Nil {
572-
tpl, err = client.Template(ctx, id)
573-
if err != nil {
574-
return xerrors.Errorf("get template by ID %q: %w", template, err)
575-
}
576-
} else {
577-
// List templates in all orgs until we find a match.
578-
orgLoop:
579-
for _, orgID := range me.OrganizationIDs {
580-
tpls, err := client.TemplatesByOrganization(ctx, orgID)
581-
if err != nil {
582-
return xerrors.Errorf("list templates in org %q: %w", orgID, err)
583-
}
584-
585-
for _, t := range tpls {
586-
if t.Name == template {
587-
tpl = t
588-
break orgLoop
589-
}
590-
}
591-
}
592-
}
593-
if tpl.ID == uuid.Nil {
594-
return xerrors.Errorf("could not find template %q in any organization", template)
589+
tpl, err := parseTemplate(ctx, client, me.OrganizationIDs, template)
590+
if err != nil {
591+
return xerrors.Errorf("parse template: %w", err)
595592
}
596593

597594
cliRichParameters, err := asWorkspaceBuildParameters(parameterFlags.richParameters)
@@ -859,6 +856,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
859856
tickInterval time.Duration
860857
bytesPerTick int64
861858
ssh bool
859+
template string
862860

863861
client = &codersdk.Client{}
864862
tracingFlags = &scaletestTracingFlags{}
@@ -876,6 +874,12 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
876874
),
877875
Handler: func(inv *clibase.Invocation) error {
878876
ctx := inv.Context()
877+
878+
me, err := requireAdmin(ctx, client)
879+
if err != nil {
880+
return err
881+
}
882+
879883
reg := prometheus.NewRegistry()
880884
metrics := workspacetraffic.NewMetrics(reg, "username", "workspace_name", "agent_name")
881885

@@ -893,7 +897,14 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
893897
},
894898
}
895899

896-
workspaces, err := getScaletestWorkspaces(inv.Context(), client)
900+
if template != "" {
901+
_, err := parseTemplate(ctx, client, me.OrganizationIDs, template)
902+
if err != nil {
903+
return xerrors.Errorf("parse template: %w", err)
904+
}
905+
}
906+
907+
workspaces, err := getScaletestWorkspaces(inv.Context(), client, template)
897908
if err != nil {
898909
return err
899910
}
@@ -997,6 +1008,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
9971008
}
9981009

9991010
cmd.Options = []clibase.Option{
1011+
{
1012+
Flag: "template",
1013+
FlagShorthand: "t",
1014+
Env: "CODER_SCALETEST_TEMPLATE",
1015+
Description: "Name or ID of the template. Traffic generation will be limited to workspaces created from this template.",
1016+
Value: clibase.StringOf(&template),
1017+
},
10001018
{
10011019
Flag: "bytes-per-tick",
10021020
Env: "CODER_SCALETEST_WORKSPACE_TRAFFIC_BYTES_PER_TICK",
@@ -1281,7 +1299,7 @@ func isScaleTestWorkspace(workspace codersdk.Workspace) bool {
12811299
strings.HasPrefix(workspace.Name, "scaletest-")
12821300
}
12831301

1284-
func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client) ([]codersdk.Workspace, error) {
1302+
func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client, template string) ([]codersdk.Workspace, error) {
12851303
var (
12861304
pageNumber = 0
12871305
limit = 100
@@ -1290,9 +1308,10 @@ func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client) ([]cod
12901308

12911309
for {
12921310
page, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{
1293-
Name: "scaletest-",
1294-
Offset: pageNumber * limit,
1295-
Limit: limit,
1311+
Name: "scaletest-",
1312+
Template: template,
1313+
Offset: pageNumber * limit,
1314+
Limit: limit,
12961315
})
12971316
if err != nil {
12981317
return nil, xerrors.Errorf("fetch scaletest workspaces page %d: %w", pageNumber, err)
@@ -1349,3 +1368,33 @@ func getScaletestUsers(ctx context.Context, client *codersdk.Client) ([]codersdk
13491368

13501369
return users, nil
13511370
}
1371+
1372+
func parseTemplate(ctx context.Context, client *codersdk.Client, organizationIDs []uuid.UUID, template string) (tpl codersdk.Template, err error) {
1373+
if id, err := uuid.Parse(template); err == nil && id != uuid.Nil {
1374+
tpl, err = client.Template(ctx, id)
1375+
if err != nil {
1376+
return tpl, xerrors.Errorf("get template by ID %q: %w", template, err)
1377+
}
1378+
} else {
1379+
// List templates in all orgs until we find a match.
1380+
orgLoop:
1381+
for _, orgID := range organizationIDs {
1382+
tpls, err := client.TemplatesByOrganization(ctx, orgID)
1383+
if err != nil {
1384+
return tpl, xerrors.Errorf("list templates in org %q: %w", orgID, err)
1385+
}
1386+
1387+
for _, t := range tpls {
1388+
if t.Name == template {
1389+
tpl = t
1390+
break orgLoop
1391+
}
1392+
}
1393+
}
1394+
}
1395+
if tpl.ID == uuid.Nil {
1396+
return tpl, xerrors.Errorf("could not find template %q in any organization", template)
1397+
}
1398+
1399+
return tpl, nil
1400+
}

‎cli/exp_scaletest_test.go

+50
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,56 @@ func TestScaleTestWorkspaceTraffic(t *testing.T) {
9191
require.ErrorContains(t, err, "no scaletest workspaces exist")
9292
}
9393

94+
// This test just validates that the CLI command accepts its known arguments.
95+
func TestScaleTestWorkspaceTraffic_Template(t *testing.T) {
96+
t.Parallel()
97+
98+
ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium)
99+
defer cancelFunc()
100+
101+
log := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
102+
client := coderdtest.New(t, &coderdtest.Options{
103+
Logger: &log,
104+
})
105+
_ = coderdtest.CreateFirstUser(t, client)
106+
107+
inv, root := clitest.New(t, "exp", "scaletest", "workspace-traffic",
108+
"--template", "doesnotexist",
109+
)
110+
clitest.SetupConfig(t, client, root)
111+
pty := ptytest.New(t)
112+
inv.Stdout = pty.Output()
113+
inv.Stderr = pty.Output()
114+
115+
err := inv.WithContext(ctx).Run()
116+
require.ErrorContains(t, err, "could not find template \"doesnotexist\" in any organization")
117+
}
118+
119+
// This test just validates that the CLI command accepts its known arguments.
120+
func TestScaleTestCleanup_Template(t *testing.T) {
121+
t.Parallel()
122+
123+
ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium)
124+
defer cancelFunc()
125+
126+
log := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
127+
client := coderdtest.New(t, &coderdtest.Options{
128+
Logger: &log,
129+
})
130+
_ = coderdtest.CreateFirstUser(t, client)
131+
132+
inv, root := clitest.New(t, "exp", "scaletest", "cleanup",
133+
"--template", "doesnotexist",
134+
)
135+
clitest.SetupConfig(t, client, root)
136+
pty := ptytest.New(t)
137+
inv.Stdout = pty.Output()
138+
inv.Stderr = pty.Output()
139+
140+
err := inv.WithContext(ctx).Run()
141+
require.ErrorContains(t, err, "could not find template \"doesnotexist\" in any organization")
142+
}
143+
94144
// This test just validates that the CLI command accepts its known arguments.
95145
func TestScaleTestDashboard(t *testing.T) {
96146
t.Parallel()

0 commit comments

Comments
 (0)