Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 43a8674

Browse files
mafredrimafredri
authored
feat(cli): add template filter support to exp scaletest cleanup and traffic (coder#10558)
1 parent 1dd3eb6 commit 43a8674

File tree

2 files changed

+130
-32
lines changed

2 files changed

+130
-32
lines changed

‎cli/exp_scaletest.go

+80-32
Original file line numberDiff line numberDiff line change
@@ -394,6 +394,8 @@ func (r *userCleanupRunner) Run(ctx context.Context, _ string, _ io.Writer) erro
394394
}
395395

396396
func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
397+
var template string
398+
397399
cleanupStrategy := &scaletestStrategyFlags{cleanup: true}
398400
client := new(codersdk.Client)
399401

@@ -407,7 +409,7 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
407409
Handler: func(inv *clibase.Invocation) error {
408410
ctx := inv.Context()
409411

410-
_, err := requireAdmin(ctx, client)
412+
me, err := requireAdmin(ctx, client)
411413
if err != nil {
412414
return err
413415
}
@@ -421,8 +423,15 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
421423
},
422424
}
423425

426+
if template != "" {
427+
_, err := parseTemplate(ctx, client, me.OrganizationIDs, template)
428+
if err != nil {
429+
return xerrors.Errorf("parse template: %w", err)
430+
}
431+
}
432+
424433
cliui.Infof(inv.Stdout, "Fetching scaletest workspaces...")
425-
workspaces, err := getScaletestWorkspaces(ctx, client)
434+
workspaces, err := getScaletestWorkspaces(ctx, client, template)
426435
if err != nil {
427436
return err
428437
}
@@ -494,6 +503,15 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
494503
},
495504
}
496505

506+
cmd.Options = clibase.OptionSet{
507+
{
508+
Flag: "template",
509+
Env: "CODER_SCALETEST_CLEANUP_TEMPLATE",
510+
Description: "Name or ID of the template. Only delete workspaces created from the given template.",
511+
Value: clibase.StringOf(&template),
512+
},
513+
}
514+
497515
cleanupStrategy.attach(&cmd.Options)
498516
return cmd
499517
}
@@ -564,34 +582,12 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
564582
return xerrors.Errorf("could not parse --output flags")
565583
}
566584

567-
var tpl codersdk.Template
568585
if template == "" {
569586
return xerrors.Errorf("--template is required")
570587
}
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)
588+
tpl, err := parseTemplate(ctx, client, me.OrganizationIDs, template)
589+
if err != nil {
590+
return xerrors.Errorf("parse template: %w", err)
595591
}
596592

597593
cliRichParameters, err := asWorkspaceBuildParameters(parameterFlags.richParameters)
@@ -859,6 +855,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
859855
tickInterval time.Duration
860856
bytesPerTick int64
861857
ssh bool
858+
template string
862859

863860
client = &codersdk.Client{}
864861
tracingFlags = &scaletestTracingFlags{}
@@ -876,6 +873,12 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
876873
),
877874
Handler: func(inv *clibase.Invocation) error {
878875
ctx := inv.Context()
876+
877+
me, err := requireAdmin(ctx, client)
878+
if err != nil {
879+
return err
880+
}
881+
879882
reg := prometheus.NewRegistry()
880883
metrics := workspacetraffic.NewMetrics(reg, "username", "workspace_name", "agent_name")
881884

@@ -893,7 +896,14 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
893896
},
894897
}
895898

896-
workspaces, err := getScaletestWorkspaces(inv.Context(), client)
899+
if template != "" {
900+
_, err := parseTemplate(ctx, client, me.OrganizationIDs, template)
901+
if err != nil {
902+
return xerrors.Errorf("parse template: %w", err)
903+
}
904+
}
905+
906+
workspaces, err := getScaletestWorkspaces(inv.Context(), client, template)
897907
if err != nil {
898908
return err
899909
}
@@ -997,6 +1007,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
9971007
}
9981008

9991009
cmd.Options = []clibase.Option{
1010+
{
1011+
Flag: "template",
1012+
FlagShorthand: "t",
1013+
Env: "CODER_SCALETEST_TEMPLATE",
1014+
Description: "Name or ID of the template. Traffic generation will be limited to workspaces created from this template.",
1015+
Value: clibase.StringOf(&template),
1016+
},
10001017
{
10011018
Flag: "bytes-per-tick",
10021019
Env: "CODER_SCALETEST_WORKSPACE_TRAFFIC_BYTES_PER_TICK",
@@ -1281,7 +1298,7 @@ func isScaleTestWorkspace(workspace codersdk.Workspace) bool {
12811298
strings.HasPrefix(workspace.Name, "scaletest-")
12821299
}
12831300

1284-
func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client) ([]codersdk.Workspace, error) {
1301+
func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client, template string) ([]codersdk.Workspace, error) {
12851302
var (
12861303
pageNumber = 0
12871304
limit = 100
@@ -1290,9 +1307,10 @@ func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client) ([]cod
12901307

12911308
for {
12921309
page, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{
1293-
Name: "scaletest-",
1294-
Offset: pageNumber * limit,
1295-
Limit: limit,
1310+
Name: "scaletest-",
1311+
Template: template,
1312+
Offset: pageNumber * limit,
1313+
Limit: limit,
12961314
})
12971315
if err != nil {
12981316
return nil, xerrors.Errorf("fetch scaletest workspaces page %d: %w", pageNumber, err)
@@ -1349,3 +1367,33 @@ func getScaletestUsers(ctx context.Context, client *codersdk.Client) ([]codersdk
13491367

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

‎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)