coder · deansheather · Nov 2, 2022 · Nov 1, 2022 · Nov 1, 2022 · Nov 2, 2022
diff --git a/cli/loadtest.go b/cli/loadtest.go
@@ -0,0 +1,179 @@
+package cli
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+	"time"
+
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/cli/cliflag"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/loadtest/harness"
+)
+
+func loadtest() *cobra.Command {
+	var (
+		configPath string
+	)
+	cmd := &cobra.Command{
+		Use:   "loadtest --config <path>",
+		Short: "Load test the Coder API",
+		// TODO: documentation and a JSON scheme file
+		Long: "Perform load tests against the Coder server. The load tests " +
+			"configurable via a JSON file.",
+		Hidden: true,
+		Args:   cobra.ExactArgs(0),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if configPath == "" {
+				return xerrors.New("config is required")
+			}
+
+			var (
+				configReader io.ReadCloser
+			)
+			if configPath == "-" {
+				configReader = io.NopCloser(cmd.InOrStdin())
+			} else {
+				f, err := os.Open(configPath)
+				if err != nil {
+					return xerrors.Errorf("open config file %q: %w", configPath, err)
+				}
+				configReader = f
+			}
+
+			var config LoadTestConfig
+			err := json.NewDecoder(configReader).Decode(&config)
+			_ = configReader.Close()
+			if err != nil {
+				return xerrors.Errorf("read config file %q: %w", configPath, err)
+			}
+
+			err = config.Validate()
+			if err != nil {
+				return xerrors.Errorf("validate config: %w", err)
+			}
+
+			client, err := CreateClient(cmd)
+			if err != nil {
+				return err
+			}
+
+			me, err := client.User(cmd.Context(), codersdk.Me)
+			if err != nil {
+				return xerrors.Errorf("fetch current user: %w", err)
+			}
+
+			// Only owners can do loadtests. This isn't a very strong check but
+			// there's not much else we can do. Ratelimits are enforced for
+			// non-owners so hopefully that limits the damage if someone
+			// disables this check and runs it against a non-owner account.
+			ok := false
+			for _, role := range me.Roles {
+				if role.Name == "owner" {
+					ok = true
+					break
+				}
+			}
+			if !ok {
+				return xerrors.Errorf("Not logged in as site owner. Load testing is only available to site owners.")
+			}
+
+			// Disable ratelimits for future requests.
+			client.BypassRatelimits = true
+
+			// Prepare the test.
+			strategy := config.Strategy.ExecutionStrategy()
+			th := harness.NewTestHarness(strategy)
+
+			for i, t := range config.Tests {
+				name := fmt.Sprintf("%s-%d", t.Type, i)
+
+				for j := 0; j < t.Count; j++ {
+					id := strconv.Itoa(j)
+					runner, err := t.NewRunner(client)
+					if err != nil {
+						return xerrors.Errorf("create %q runner for %s/%s: %w", t.Type, name, id, err)
+					}
+
+					th.AddRun(name, id, runner)
+				}
+			}
+
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "Running load test...")
+
+			testCtx := cmd.Context()
+			if config.Timeout > 0 {
+				var cancel func()
+				testCtx, cancel = context.WithTimeout(testCtx, time.Duration(config.Timeout))
+				defer cancel()
+			}
+
+			// TODO: live progress output
+			start := time.Now()
+			err = th.Run(testCtx)
+			if err != nil {
+				return xerrors.Errorf("run test harness (harness failure, not a test failure): %w", err)
+			}
+			elapsed := time.Since(start)
+
+			// Print the results.
+			// TODO: better result printing
+			// TODO: move result printing to the loadtest package, add multiple
+			//       output formats (like HTML, JSON)
+			res := th.Results()
+			var totalDuration time.Duration
+			for _, run := range res.Runs {
+				totalDuration += run.Duration
+				if run.Error == nil {
+					continue
+				}
+
+				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\n== FAIL: %s\n\n", run.FullID)
+				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tError: %s\n\n", run.Error)
+
+				// Print log lines indented.
+				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tLog:\n")
+				rd := bufio.NewReader(bytes.NewBuffer(run.Logs))
+				for {
+					line, err := rd.ReadBytes('\n')
+					if err == io.EOF {
+						break
+					}
+					if err != nil {
+						_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\n\tLOG PRINT ERROR: %+v\n", err)
+					}
+
+					_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\t\t%s", line)
+				}
+			}
+
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\n\nTest results:")
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tPass:  %d\n", res.TotalPass)
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tFail:  %d\n", res.TotalFail)
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tTotal: %d\n", res.TotalRuns)
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "")
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tTotal duration: %s\n", elapsed)
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tAvg. duration:  %s\n", totalDuration/time.Duration(res.TotalRuns))
+
+			// Cleanup.
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\nCleaning up...")
+			err = th.Cleanup(cmd.Context())
+			if err != nil {
+				return xerrors.Errorf("cleanup tests: %w", err)
+			}
+
+			return nil
+		},
+	}
+
+	cliflag.StringVarP(cmd.Flags(), &configPath, "config", "", "CODER_LOADTEST_CONFIG_PATH", "", "Path to the load test configuration file, or - to read from stdin.")
+	return cmd
+}
diff --git a/cli/loadtest_test.go b/cli/loadtest_test.go
@@ -0,0 +1,136 @@
+package cli_test
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cli"
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/coderd/httpapi"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/loadtest/placebo"
+	"github.com/coder/coder/loadtest/workspacebuild"
+	"github.com/coder/coder/pty/ptytest"
+	"github.com/coder/coder/testutil"
+)
+
+func TestLoadTest(t *testing.T) {
+	t.Parallel()
+
+	t.Run("PlaceboFromStdin", func(t *testing.T) {
+		t.Parallel()
+
+		client := coderdtest.New(t, nil)
+		_ = coderdtest.CreateFirstUser(t, client)
+
+		config := cli.LoadTestConfig{
+			Strategy: cli.LoadTestStrategy{
+				Type: cli.LoadTestStrategyTypeLinear,
+			},
+			Tests: []cli.LoadTest{
+				{
+					Type:  cli.LoadTestTypePlacebo,
+					Count: 10,
+					Placebo: &placebo.Config{
+						Sleep: httpapi.Duration(10 * time.Millisecond),
+					},
+				},
+			},
+			Timeout: httpapi.Duration(testutil.WaitShort),
+		}
+
+		configBytes, err := json.Marshal(config)
+		require.NoError(t, err)
+
+		cmd, root := clitest.New(t, "loadtest", "--config", "-")
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t)
+		cmd.SetIn(bytes.NewReader(configBytes))
+		cmd.SetOut(pty.Output())
+		cmd.SetErr(pty.Output())
+
+		ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancelFunc()
+
+		done := make(chan any)
+		go func() {
+			errC := cmd.ExecuteContext(ctx)
+			assert.NoError(t, errC)
+			close(done)
+		}()
+		pty.ExpectMatch("Test results:")
+		pty.ExpectMatch("Pass:  10")
+		cancelFunc()
+		<-done
+	})
+
+	t.Run("WorkspaceBuildFromFile", func(t *testing.T) {
+		t.Parallel()
+
+		client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+		user := coderdtest.CreateFirstUser(t, client)
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+
+		config := cli.LoadTestConfig{
+			Strategy: cli.LoadTestStrategy{
+				Type:             cli.LoadTestStrategyTypeConcurrent,
+				ConcurrencyLimit: 2,
+			},
+			Tests: []cli.LoadTest{
+				{
+					Type:  cli.LoadTestTypeWorkspaceBuild,
+					Count: 2,
+					WorkspaceBuild: &workspacebuild.Config{
+						OrganizationID: user.OrganizationID,
+						UserID:         user.UserID.String(),
+						Request: codersdk.CreateWorkspaceRequest{
+							TemplateID: template.ID,
+						},
+					},
+				},
+			},
+			Timeout: httpapi.Duration(testutil.WaitLong),
+		}
+
+		d := t.TempDir()
+		configPath := filepath.Join(d, "/config.loadtest.json")
+		f, err := os.Create(configPath)
+		require.NoError(t, err)
+		defer f.Close()
+		err = json.NewEncoder(f).Encode(config)
+		require.NoError(t, err)
+		_ = f.Close()
+
+		cmd, root := clitest.New(t, "loadtest", "--config", configPath)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
+		cmd.SetErr(pty.Output())
+
+		ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancelFunc()
+
+		done := make(chan any)
+		go func() {
+			errC := cmd.ExecuteContext(ctx)
+			assert.NoError(t, errC)
+			close(done)
+		}()
+		pty.ExpectMatch("Test results:")
+		pty.ExpectMatch("Pass:  2")
+		<-done
+		cancelFunc()
+	})
+}