Skip to content

Commit 8c48d2b

Browse files
committed
move test to workspacetraffic package
1 parent d54af33 commit 8c48d2b

File tree

2 files changed

+121
-54
lines changed

2 files changed

+121
-54
lines changed

cli/scaletest_test.go

Lines changed: 6 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,12 @@ import (
88
"path/filepath"
99
"testing"
1010

11-
"github.com/google/uuid"
1211
"github.com/stretchr/testify/assert"
1312
"github.com/stretchr/testify/require"
1413

15-
"github.com/coder/coder/agent"
1614
"github.com/coder/coder/cli/clitest"
1715
"github.com/coder/coder/coderd/coderdtest"
1816
"github.com/coder/coder/codersdk"
19-
"github.com/coder/coder/codersdk/agentsdk"
20-
"github.com/coder/coder/provisioner/echo"
21-
"github.com/coder/coder/provisionersdk/proto"
2217
"github.com/coder/coder/pty/ptytest"
2318
"github.com/coder/coder/scaletest/harness"
2419
"github.com/coder/coder/testutil"
@@ -205,70 +200,27 @@ param3: 1
205200
})
206201
}
207202

208-
// This test pretends to stand up a workspace and run a no-op traffic generation test.
209-
// It's not a real test, but it's useful for debugging.
210-
// We do not perform any cleanup.
203+
// This test just validates that the CLI command accepts its known arguments.
204+
// A more comprehensive test is performed in workspacetraffic/run_test.go
211205
func TestScaleTestWorkspaceTraffic(t *testing.T) {
212206
t.Parallel()
213207

214208
ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium)
215209
defer cancelFunc()
216210

217-
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
218-
user := coderdtest.CreateFirstUser(t, client)
219-
220-
authToken := uuid.NewString()
221-
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
222-
Parse: echo.ParseComplete,
223-
ProvisionPlan: echo.ProvisionComplete,
224-
ProvisionApply: []*proto.Provision_Response{{
225-
Type: &proto.Provision_Response_Complete{
226-
Complete: &proto.Provision_Complete{
227-
Resources: []*proto.Resource{{
228-
Name: "example",
229-
Type: "aws_instance",
230-
Agents: []*proto.Agent{{
231-
Id: uuid.NewString(),
232-
Name: "agent",
233-
Auth: &proto.Agent_Token{
234-
Token: authToken,
235-
},
236-
Apps: []*proto.App{},
237-
}},
238-
}},
239-
},
240-
},
241-
}},
242-
})
243-
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
244-
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
245-
246-
ws := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
247-
cwr.Name = "scaletest-test"
248-
})
249-
coderdtest.AwaitWorkspaceBuildJob(t, client, ws.LatestBuild.ID)
250-
251-
agentClient := agentsdk.New(client.URL)
252-
agentClient.SetSessionToken(authToken)
253-
agentCloser := agent.New(agent.Options{
254-
Client: agentClient,
255-
})
256-
t.Cleanup(func() {
257-
_ = agentCloser.Close()
258-
})
259-
260-
coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
211+
client := coderdtest.New(t, nil)
212+
_ = coderdtest.CreateFirstUser(t, client)
261213

262214
inv, root := clitest.New(t, "scaletest", "workspace-traffic",
263215
"--timeout", "1s",
264216
"--bytes-per-tick", "1024",
265217
"--tick-interval", "100ms",
218+
"--scaletest-prometheus-address", "127.0.0.1:0",
266219
)
267220
clitest.SetupConfig(t, client, root)
268221
var stdout, stderr bytes.Buffer
269222
inv.Stdout = &stdout
270223
inv.Stderr = &stderr
271224
err := inv.WithContext(ctx).Run()
272-
require.NoError(t, err)
273-
require.Contains(t, stdout.String(), "Pass: 1")
225+
require.ErrorContains(t, err, "no scaletest workspaces exist")
274226
}
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package workspacetraffic_test
2+
3+
import (
4+
"context"
5+
"strings"
6+
"testing"
7+
"time"
8+
9+
"github.com/coder/coder/agent"
10+
"github.com/coder/coder/coderd/coderdtest"
11+
"github.com/coder/coder/codersdk"
12+
"github.com/coder/coder/codersdk/agentsdk"
13+
"github.com/coder/coder/provisioner/echo"
14+
"github.com/coder/coder/provisionersdk/proto"
15+
"github.com/coder/coder/scaletest/workspacetraffic"
16+
"github.com/coder/coder/testutil"
17+
18+
"github.com/google/uuid"
19+
"github.com/prometheus/client_golang/prometheus"
20+
"github.com/stretchr/testify/assert"
21+
"github.com/stretchr/testify/require"
22+
)
23+
24+
func TestRun(t *testing.T) {
25+
t.Parallel()
26+
27+
// We need to stand up an in-memory coderd and run a fake workspace.
28+
var (
29+
client = coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
30+
user = coderdtest.CreateFirstUser(t, client)
31+
authToken = uuid.NewString()
32+
agentName = "agent"
33+
version = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
34+
Parse: echo.ParseComplete,
35+
ProvisionPlan: echo.ProvisionComplete,
36+
ProvisionApply: []*proto.Provision_Response{{
37+
Type: &proto.Provision_Response_Complete{
38+
Complete: &proto.Provision_Complete{
39+
Resources: []*proto.Resource{{
40+
Name: "example",
41+
Type: "aws_instance",
42+
Agents: []*proto.Agent{{
43+
// Agent ID gets generated no matter what we say ¯\_(ツ)_/¯
44+
Name: agentName,
45+
Auth: &proto.Agent_Token{
46+
Token: authToken,
47+
},
48+
Apps: []*proto.App{},
49+
}},
50+
}},
51+
},
52+
},
53+
}},
54+
})
55+
template = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
56+
_ = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
57+
// In order to be picked up as a scaletest workspace, the workspace must be named specifically
58+
ws = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
59+
cwr.Name = "scaletest-test"
60+
})
61+
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, ws.LatestBuild.ID)
62+
)
63+
64+
// We also need a running agent to run this test.
65+
agentClient := agentsdk.New(client.URL)
66+
agentClient.SetSessionToken(authToken)
67+
agentCloser := agent.New(agent.Options{
68+
Client: agentClient,
69+
})
70+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
71+
72+
t.Cleanup(cancel)
73+
t.Cleanup(func() {
74+
_ = agentCloser.Close()
75+
})
76+
// Make sure the agent is connected before we go any further.
77+
resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
78+
var agentID uuid.UUID
79+
for _, res := range resources {
80+
for _, agt := range res.Agents {
81+
agentID = agt.ID
82+
}
83+
}
84+
require.NotEqual(t, uuid.Nil, agentID, "did not expect agentID to be nil")
85+
86+
// Now we can start the runner.
87+
reg := prometheus.NewRegistry()
88+
metrics := workspacetraffic.NewMetrics(reg, "username", "workspace_name", "agent_name")
89+
runner := workspacetraffic.NewRunner(client, workspacetraffic.Config{
90+
AgentID: agentID,
91+
AgentName: agentName,
92+
WorkspaceName: ws.Name,
93+
WorkspaceOwner: ws.OwnerName,
94+
BytesPerTick: 1024,
95+
TickInterval: testutil.IntervalMedium,
96+
Duration: testutil.WaitMedium - time.Second,
97+
Registry: reg,
98+
}, metrics)
99+
100+
var logs strings.Builder
101+
require.NoError(t, runner.Run(ctx, "", &logs), "unexpected error calling Run()")
102+
103+
var collected []prometheus.Metric
104+
collectCh := make(chan prometheus.Metric)
105+
go func() {
106+
for metric := range collectCh {
107+
collected = append(collected, metric)
108+
}
109+
}()
110+
reg.Collect(collectCh)
111+
assert.NotEmpty(t, collected)
112+
for _, m := range collected {
113+
assert.NotZero(t, m.Desc())
114+
}
115+
}

0 commit comments

Comments
 (0)