"fix" dashboard test

johnstcn · johnstcn · commit d65e2c83be13 · 2023-10-10T21:55:44.000Z
diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go
@@ -1151,13 +1151,12 @@ func (r *RootCmd) scaletestDashboard() *clibase.Cmd {
 				userClient.SetSessionToken(userTokResp.Key)
 
 				config := dashboard.Config{
-					Interval:   interval,
-					Jitter:     jitter,
-					Trace:      tracingEnabled,
-					Logger:     logger.Named(name),
-					Headless:   headless,
-					ActionFunc: dashboard.ClickRandomElement,
-					RandIntn:   rndGen.Intn,
+					Interval: interval,
+					Jitter:   jitter,
+					Trace:    tracingEnabled,
+					Logger:   logger.Named(name),
+					Headless: headless,
+					RandIntn: rndGen.Intn,
 				}
 				//nolint:gocritic
 				logger.Info(ctx, "runner config", slog.F("interval", interval), slog.F("jitter", jitter), slog.F("headless", headless), slog.F("trace", tracingEnabled))
diff --git a/scaletest/dashboard/chromedp.go b/scaletest/dashboard/chromedp.go
@@ -88,11 +88,11 @@ var defaultTargets = []Target{
 	},
 }
 
-// ClickRandomElement returns an action that will click an element from defaultTargets.
+// clickRandomElement returns an action that will click an element from defaultTargets.
 // If no elements are found, an error is returned.
 // If more than one element is found, one is chosen at random.
 // The label of the clicked element is returned.
-func ClickRandomElement(ctx context.Context, log slog.Logger, randIntn func(int) int, deadline time.Time) (Label, Action, error) {
+func clickRandomElement(ctx context.Context, log slog.Logger, randIntn func(int) int, deadline time.Time) (Label, Action, error) {
 	var xpath Selector
 	var found bool
 	var err error
diff --git a/scaletest/dashboard/config.go b/scaletest/dashboard/config.go
@@ -22,6 +22,10 @@ type Config struct {
 	Headless bool `json:"headless"`
 	// ActionFunc is a function that returns an action to run.
 	ActionFunc func(ctx context.Context, log slog.Logger, randIntn func(int) int, deadline time.Time) (Label, Action, error) `json:"-"`
+	// WaitLoaded is a function that waits for the page to be loaded.
+	WaitLoaded func(ctx context.Context, deadline time.Time) error
+	// Screenshot is a function that takes a screenshot.
+	Screenshot func(ctx context.Context, filename string) (string, error)
 	// RandIntn is a function that returns a random number between 0 and n-1.
 	RandIntn func(int) int `json:"-"`
 }
diff --git a/scaletest/dashboard/run.go b/scaletest/dashboard/run.go
@@ -2,6 +2,7 @@ package dashboard
 
 import (
 	"context"
+	"errors"
 	"io"
 	"time"
 
@@ -25,6 +26,15 @@ var (
 
 func NewRunner(client *codersdk.Client, metrics Metrics, cfg Config) *Runner {
 	client.Trace = cfg.Trace
+	if cfg.WaitLoaded == nil {
+		cfg.WaitLoaded = waitForWorkspacesPageLoaded
+	}
+	if cfg.ActionFunc == nil {
+		cfg.ActionFunc = clickRandomElement
+	}
+	if cfg.Screenshot == nil {
+		cfg.Screenshot = screenshot
+	}
 	return &Runner{
 		client:  client,
 		cfg:     cfg,
@@ -33,6 +43,16 @@ func NewRunner(client *codersdk.Client, metrics Metrics, cfg Config) *Runner {
 }
 
 func (r *Runner) Run(ctx context.Context, _ string, _ io.Writer) error {
+	err := r._run(ctx)
+	// If the context deadline exceeded, don't return an error.
+	// This just means the test finished.
+	if err == nil || errors.Is(err, context.DeadlineExceeded) {
+		return nil
+	}
+	return err
+}
+
+func (r *Runner) _run(ctx context.Context) error {
 	if r.client == nil {
 		return xerrors.Errorf("client is nil")
 	}
@@ -55,7 +75,7 @@ func (r *Runner) Run(ctx context.Context, _ string, _ io.Writer) error {
 	defer t.Stop()
 	r.cfg.Logger.Info(ctx, "waiting for workspaces page to load")
 	loadWorkspacePageDeadline := time.Now().Add(r.cfg.Interval)
-	if err := waitForWorkspacesPageLoaded(cdpCtx, loadWorkspacePageDeadline); err != nil {
+	if err := r.cfg.WaitLoaded(cdpCtx, loadWorkspacePageDeadline); err != nil {
 		return xerrors.Errorf("wait for workspaces page to load: %w", err)
 	}
 	for {
@@ -73,7 +93,7 @@ func (r *Runner) Run(ctx context.Context, _ string, _ io.Writer) error {
 			l, act, err := r.cfg.ActionFunc(cdpCtx, r.cfg.Logger, r.cfg.RandIntn, actionCompleteByDeadline)
 			if err != nil {
 				r.cfg.Logger.Error(ctx, "calling ActionFunc", slog.Error(err))
-				sPath, sErr := screenshot(cdpCtx, me.Username)
+				sPath, sErr := r.cfg.Screenshot(cdpCtx, me.Username)
 				if sErr != nil {
 					r.cfg.Logger.Error(ctx, "screenshot failed", slog.Error(sErr))
 				}
@@ -88,7 +108,7 @@ func (r *Runner) Run(ctx context.Context, _ string, _ io.Writer) error {
 				r.metrics.IncErrors(string(l))
 				//nolint:gocritic
 				r.cfg.Logger.Error(ctx, "action failed", slog.F("label", l), slog.Error(err))
-				sPath, sErr := screenshot(cdpCtx, me.Username+"-"+string(l))
+				sPath, sErr := r.cfg.Screenshot(cdpCtx, me.Username+"-"+string(l))
 				if sErr != nil {
 					r.cfg.Logger.Error(ctx, "screenshot failed", slog.Error(sErr))
 				}
diff --git a/scaletest/dashboard/run_test.go b/scaletest/dashboard/run_test.go
@@ -5,6 +5,7 @@ import (
 	"math/rand"
 	"runtime"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -47,17 +48,29 @@ func Test_Run(t *testing.T) {
 		IgnoreErrors: true,
 	})
 	m := &testMetrics{}
+	var (
+		waitLoadedCalled atomic.Bool
+		screenshotCalled atomic.Bool
+	)
 	cfg := dashboard.Config{
 		Interval: 500 * time.Millisecond,
 		Jitter:   100 * time.Millisecond,
 		Logger:   log,
 		Headless: true,
+		WaitLoaded: func(_ context.Context, _ time.Time) error {
+			waitLoadedCalled.Store(true)
+			return nil
+		},
 		ActionFunc: func(_ context.Context, _ slog.Logger, rnd func(int) int, _ time.Time) (dashboard.Label, dashboard.Action, error) {
 			if rnd(2) == 0 {
 				return "fails", failAction, nil
 			}
 			return "succeeds", successAction, nil
 		},
+		Screenshot: func(_ context.Context, name string) (string, error) {
+			screenshotCalled.Store(true)
+			return "/fake/path/to/" + name + ".png", nil
+		},
 		RandIntn: rg.Intn,
 	}
 	r := dashboard.NewRunner(client, m, cfg)
