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

Skip to content

Commit 4839db7

Browse files
deansheatherdeansheather
committed
feat: add coder loadtest command
1 parent b63d7e8 commit 4839db7

File tree

15 files changed

+906
-29
lines changed

15 files changed

+906
-29
lines changed

‎cli/loadtest.go

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
package cli
2+
3+
import (
4+
"bufio"
5+
"bytes"
6+
"context"
7+
"encoding/json"
8+
"fmt"
9+
"io"
10+
"os"
11+
"strconv"
12+
"time"
13+
14+
"github.com/spf13/cobra"
15+
"golang.org/x/xerrors"
16+
17+
"github.com/coder/coder/cli/cliflag"
18+
"github.com/coder/coder/codersdk"
19+
"github.com/coder/coder/loadtest/harness"
20+
)
21+
22+
func loadtest() *cobra.Command {
23+
var (
24+
configPath string
25+
)
26+
cmd := &cobra.Command{
27+
Use: "loadtest --config <path>",
28+
Short: "Load test the Coder API",
29+
Long: "",
30+
Args: cobra.ExactArgs(0),
31+
RunE: func(cmd *cobra.Command, args []string) error {
32+
if configPath == "" {
33+
return xerrors.New("config is required")
34+
}
35+
36+
var (
37+
configReader io.Reader
38+
configCloser io.Closer
39+
)
40+
if configPath == "-" {
41+
configReader = cmd.InOrStdin()
42+
} else {
43+
f, err := os.Open(configPath)
44+
if err != nil {
45+
return xerrors.Errorf("open config file %q: %w", configPath, err)
46+
}
47+
configReader = f
48+
configCloser = f
49+
}
50+
51+
var config LoadTestConfig
52+
err := json.NewDecoder(configReader).Decode(&config)
53+
if configCloser != nil {
54+
_ = configCloser.Close()
55+
}
56+
if err != nil {
57+
return xerrors.Errorf("read config file %q: %w", configPath, err)
58+
}
59+
60+
err = config.Validate()
61+
if err != nil {
62+
return xerrors.Errorf("validate config: %w", err)
63+
}
64+
65+
client, err := CreateClient(cmd)
66+
if err != nil {
67+
return err
68+
}
69+
70+
me, err := client.User(cmd.Context(), codersdk.Me)
71+
if err != nil {
72+
return xerrors.Errorf("fetch current user: %w", err)
73+
}
74+
75+
// Only owners can do loadtests. This isn't a very strong check but
76+
// there's not much else we can do. Ratelimits are enforced for
77+
// non-owners so hopefully that limits the damage if someone
78+
// disables this check and runs it against a non-owner account.
79+
ok := false
80+
for _, role := range me.Roles {
81+
if role.Name == "owner" {
82+
ok = true
83+
break
84+
}
85+
}
86+
if !ok {
87+
return xerrors.Errorf("Not logged in as site owner. Load testing is only available to the site owner.")
88+
}
89+
90+
// Disable ratelimits for future requests.
91+
client.BypassRatelimits = true
92+
93+
// Prepare the test.
94+
strategy := config.Strategy.ExecutionStrategy()
95+
th := harness.NewTestHarness(strategy)
96+
97+
for i, t := range config.Tests {
98+
name := fmt.Sprintf("%s-%d", t.Type, i)
99+
100+
for i := 0; i < t.Count; i++ {
101+
id := strconv.Itoa(i)
102+
runner, err := t.NewRunner(client)
103+
if err != nil {
104+
return xerrors.Errorf("create %q runner for %s/%s: %w", t.Type, name, id, err)
105+
}
106+
107+
th.AddRun(name, id, runner)
108+
}
109+
}
110+
111+
_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "Running load test...")
112+
113+
testCtx := cmd.Context()
114+
if config.Timeout > 0 {
115+
var cancel func()
116+
testCtx, cancel = context.WithTimeout(testCtx, config.Timeout)
117+
defer cancel()
118+
}
119+
120+
// TODO: live progress output
121+
start := time.Now()
122+
err = th.Run(testCtx)
123+
if err != nil {
124+
return xerrors.Errorf("run test harness (harness failure, not a test failure): %w", err)
125+
}
126+
elapsed := time.Since(start)
127+
128+
// Print the results.
129+
// TODO: better result printing
130+
// TODO: move result printing to the loadtest package, add multiple
131+
// output formats (like HTML, JSON)
132+
res := th.Results()
133+
var totalDuration time.Duration
134+
for _, run := range res.Runs {
135+
totalDuration += run.Duration
136+
if run.Error == nil {
137+
continue
138+
}
139+
140+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\n== FAIL: %s\n\n", run.FullID)
141+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tError: %s\n\n", run.Error)
142+
143+
// Print log lines indented.
144+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tLog:\n")
145+
rd := bufio.NewReader(bytes.NewBuffer(run.Logs))
146+
for {
147+
line, err := rd.ReadBytes('\n')
148+
if err == io.EOF {
149+
break
150+
}
151+
if err != nil {
152+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\n\tLOG PRINT ERROR: %+v\n", err)
153+
}
154+
155+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\t\t%s", line)
156+
}
157+
}
158+
159+
_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\n\nTest results:")
160+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tPass: %d\n", res.TotalPass)
161+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tFail: %d\n", res.TotalFail)
162+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tTotal: %d\n", res.TotalRuns)
163+
_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "")
164+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tTotal duration: %s\n", elapsed)
165+
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tAvg. duration: %s\n", totalDuration/time.Duration(res.TotalRuns))
166+
167+
// Cleanup.
168+
_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\nCleaning up...")
169+
err = th.Cleanup(cmd.Context())
170+
if err != nil {
171+
return xerrors.Errorf("cleanup tests: %w", err)
172+
}
173+
174+
return nil
175+
},
176+
}
177+
178+
cliflag.StringVarP(cmd.Flags(), &configPath, "config", "", "CODER_LOADTEST_CONFIG_PATH", "", "Path to the load test configuration file, or - to read from stdin.")
179+
return cmd
180+
}

‎cli/loadtest_test.go

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
package cli_test
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"encoding/json"
7+
"os"
8+
"path/filepath"
9+
"testing"
10+
"time"
11+
12+
"github.com/stretchr/testify/assert"
13+
"github.com/stretchr/testify/require"
14+
15+
"github.com/coder/coder/cli"
16+
"github.com/coder/coder/cli/clitest"
17+
"github.com/coder/coder/coderd/coderdtest"
18+
"github.com/coder/coder/codersdk"
19+
"github.com/coder/coder/loadtest/placebo"
20+
"github.com/coder/coder/loadtest/workspacebuild"
21+
"github.com/coder/coder/pty/ptytest"
22+
"github.com/coder/coder/testutil"
23+
)
24+
25+
func TestLoadTest(t *testing.T) {
26+
t.Parallel()
27+
28+
t.Run("PlaceboFromStdin", func(t *testing.T) {
29+
t.Parallel()
30+
31+
client := coderdtest.New(t, nil)
32+
_ = coderdtest.CreateFirstUser(t, client)
33+
34+
config := cli.LoadTestConfig{
35+
Strategy: cli.LoadTestStrategy{
36+
Type: cli.LoadTestStrategyTypeLinear,
37+
},
38+
Tests: []cli.LoadTest{
39+
{
40+
Type: cli.LoadTestTypePlacebo,
41+
Count: 10,
42+
Placebo: &placebo.Config{
43+
Sleep: 10 * time.Millisecond,
44+
},
45+
},
46+
},
47+
Timeout: 1 * time.Second,
48+
}
49+
50+
configBytes, err := json.Marshal(config)
51+
require.NoError(t, err)
52+
53+
cmd, root := clitest.New(t, "loadtest", "--config", "-")
54+
clitest.SetupConfig(t, client, root)
55+
pty := ptytest.New(t)
56+
cmd.SetIn(bytes.NewReader(configBytes))
57+
cmd.SetOut(pty.Output())
58+
cmd.SetErr(pty.Output())
59+
60+
ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitLong)
61+
defer cancelFunc()
62+
63+
done := make(chan any)
64+
go func() {
65+
errC := cmd.ExecuteContext(ctx)
66+
assert.NoError(t, errC)
67+
close(done)
68+
}()
69+
pty.ExpectMatch("Test results:")
70+
pty.ExpectMatch("Pass: 10")
71+
cancelFunc()
72+
<-done
73+
})
74+
75+
t.Run("WorkspaceBuildFromFile", func(t *testing.T) {
76+
t.Parallel()
77+
78+
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
79+
user := coderdtest.CreateFirstUser(t, client)
80+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
81+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
82+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
83+
84+
config := cli.LoadTestConfig{
85+
Strategy: cli.LoadTestStrategy{
86+
Type: cli.LoadTestStrategyTypeConcurrent,
87+
ConcurrencyLimit: 2,
88+
},
89+
Tests: []cli.LoadTest{
90+
{
91+
Type: cli.LoadTestTypeWorkspaceBuild,
92+
Count: 2,
93+
WorkspaceBuild: &workspacebuild.Config{
94+
OrganizationID: user.OrganizationID,
95+
UserID: user.UserID.String(),
96+
Request: codersdk.CreateWorkspaceRequest{
97+
TemplateID: template.ID,
98+
},
99+
},
100+
},
101+
},
102+
Timeout: 10 * time.Second,
103+
}
104+
105+
d := t.TempDir()
106+
configPath := filepath.Join(d, "/config.loadtest.json")
107+
f, err := os.Create(configPath)
108+
require.NoError(t, err)
109+
defer f.Close()
110+
err = json.NewEncoder(f).Encode(config)
111+
require.NoError(t, err)
112+
_ = f.Close()
113+
114+
cmd, root := clitest.New(t, "loadtest", "--config", configPath)
115+
clitest.SetupConfig(t, client, root)
116+
pty := ptytest.New(t)
117+
cmd.SetIn(pty.Input())
118+
cmd.SetOut(pty.Output())
119+
cmd.SetErr(pty.Output())
120+
121+
ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitLong)
122+
defer cancelFunc()
123+
124+
done := make(chan any)
125+
go func() {
126+
errC := cmd.ExecuteContext(ctx)
127+
assert.NoError(t, errC)
128+
close(done)
129+
}()
130+
pty.ExpectMatch("Test results:")
131+
pty.ExpectMatch("Pass: 2")
132+
<-done
133+
cancelFunc()
134+
})
135+
}

0 commit comments

Comments
 (0)