Skip to content

Commit e847276

Browse files
deansheatherdeansheather
authored
feat: add cleanup strategy to loadtest (#4991)
1 parent 1c9677d commit e847276

9 files changed

+314
-122
lines changed

‎cli/loadtest.go

+3-2
Original file line numberDiff line numberDiff line change
@@ -135,8 +135,9 @@ func loadtest() *cobra.Command {
135135
client.PropagateTracing = tracePropagate
136136

137137
// Prepare the test.
138-
strategy := config.Strategy.ExecutionStrategy()
139-
th := harness.NewTestHarness(strategy)
138+
runStrategy := config.Strategy.ExecutionStrategy()
139+
cleanupStrategy := config.CleanupStrategy.ExecutionStrategy()
140+
th := harness.NewTestHarness(runStrategy, cleanupStrategy)
140141

141142
for i, t := range config.Tests {
142143
name := fmt.Sprintf("%s-%d", t.Type, i)

‎cli/loadtest_test.go

+10
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ func TestLoadTest(t *testing.T) {
3838
Strategy: cli.LoadTestStrategy{
3939
Type: cli.LoadTestStrategyTypeLinear,
4040
},
41+
CleanupStrategy: cli.LoadTestStrategy{
42+
Type: cli.LoadTestStrategyTypeLinear,
43+
},
4144
Tests: []cli.LoadTest{
4245
{
4346
Type: cli.LoadTestTypePlacebo,
@@ -89,6 +92,10 @@ func TestLoadTest(t *testing.T) {
8992
Type: cli.LoadTestStrategyTypeConcurrent,
9093
ConcurrencyLimit: 2,
9194
},
95+
CleanupStrategy: cli.LoadTestStrategy{
96+
Type: cli.LoadTestStrategyTypeConcurrent,
97+
ConcurrencyLimit: 2,
98+
},
9299
Tests: []cli.LoadTest{
93100
{
94101
Type: cli.LoadTestTypeWorkspaceBuild,
@@ -210,6 +217,9 @@ func TestLoadTest(t *testing.T) {
210217
Strategy: cli.LoadTestStrategy{
211218
Type: cli.LoadTestStrategyTypeLinear,
212219
},
220+
CleanupStrategy: cli.LoadTestStrategy{
221+
Type: cli.LoadTestStrategyTypeLinear,
222+
},
213223
Tests: []cli.LoadTest{
214224
{
215225
Type: cli.LoadTestTypePlacebo,

‎cli/loadtestconfig.go

+7-2
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@ import (
1515

1616
// LoadTestConfig is the overall configuration for a call to `coder loadtest`.
1717
type LoadTestConfig struct {
18-
Strategy LoadTestStrategy `json:"strategy"`
19-
Tests []LoadTest `json:"tests"`
18+
Strategy LoadTestStrategy `json:"strategy"`
19+
CleanupStrategy LoadTestStrategy `json:"cleanup_strategy"`
20+
Tests []LoadTest `json:"tests"`
2021
// Timeout sets a timeout for the entire test run, to control the timeout
2122
// for each individual run use strategy.timeout.
2223
Timeout httpapi.Duration `json:"timeout"`
@@ -134,6 +135,10 @@ func (c *LoadTestConfig) Validate() error {
134135
if err != nil {
135136
return xerrors.Errorf("validate strategy: %w", err)
136137
}
138+
err = c.CleanupStrategy.Validate()
139+
if err != nil {
140+
return xerrors.Errorf("validate cleanup_strategy: %w", err)
141+
}
137142

138143
for i, test := range c.Tests {
139144
err := test.Validate()

‎loadtest/harness/harness.go

+40-26
Original file line numberDiff line numberDiff line change
@@ -11,19 +11,11 @@ import (
1111
"github.com/coder/coder/coderd/tracing"
1212
)
1313

14-
// ExecutionStrategy defines how a TestHarness should execute a set of runs. It
15-
// essentially defines the concurrency model for a given testing session.
16-
type ExecutionStrategy interface {
17-
// Execute runs the given runs in whatever way the strategy wants. An error
18-
// may only be returned if the strategy has a failure itself, not if any of
19-
// the runs fail.
20-
Execute(ctx context.Context, runs []*TestRun) error
21-
}
22-
23-
// TestHarness runs a bunch of registered test runs using the given
24-
// ExecutionStrategy.
14+
// TestHarness runs a bunch of registered test runs using the given execution
15+
// strategies.
2516
type TestHarness struct {
26-
strategy ExecutionStrategy
17+
runStrategy ExecutionStrategy
18+
cleanupStrategy ExecutionStrategy
2719

2820
mut *sync.Mutex
2921
runIDs map[string]struct{}
@@ -33,14 +25,15 @@ type TestHarness struct {
3325
elapsed time.Duration
3426
}
3527

36-
// NewTestHarness creates a new TestHarness with the given ExecutionStrategy.
37-
func NewTestHarness(strategy ExecutionStrategy) *TestHarness {
28+
// NewTestHarness creates a new TestHarness with the given execution strategies.
29+
func NewTestHarness(runStrategy, cleanupStrategy ExecutionStrategy) *TestHarness {
3830
return &TestHarness{
39-
strategy: strategy,
40-
mut: new(sync.Mutex),
41-
runIDs: map[string]struct{}{},
42-
runs: []*TestRun{},
43-
done: make(chan struct{}),
31+
runStrategy: runStrategy,
32+
cleanupStrategy: cleanupStrategy,
33+
mut: new(sync.Mutex),
34+
runIDs: map[string]struct{}{},
35+
runs: []*TestRun{},
36+
done: make(chan struct{}),
4437
}
4538
}
4639

@@ -62,11 +55,16 @@ func (h *TestHarness) Run(ctx context.Context) (err error) {
6255
h.started = true
6356
h.mut.Unlock()
6457

58+
runFns := make([]TestFn, len(h.runs))
59+
for i, run := range h.runs {
60+
runFns[i] = run.Run
61+
}
62+
6563
defer close(h.done)
6664
defer func() {
6765
e := recover()
6866
if e != nil {
69-
err = xerrors.Errorf("execution strategy panicked: %w", e)
67+
err = xerrors.Errorf("panic in harness.Run: %+v", e)
7068
}
7169
}()
7270

@@ -77,7 +75,9 @@ func (h *TestHarness) Run(ctx context.Context) (err error) {
7775
h.elapsed = time.Since(start)
7876
}()
7977

80-
err = h.strategy.Execute(ctx, h.runs)
78+
// We don't care about test failures here since they already get recorded
79+
// by the *TestRun.
80+
_, err = h.runStrategy.Run(ctx, runFns)
8181
//nolint:revive // we use named returns because we mutate it in a defer
8282
return
8383
}
@@ -96,20 +96,34 @@ func (h *TestHarness) Cleanup(ctx context.Context) (err error) {
9696
panic("harness has not finished")
9797
}
9898

99+
cleanupFns := make([]TestFn, len(h.runs))
100+
for i, run := range h.runs {
101+
cleanupFns[i] = run.Cleanup
102+
}
103+
99104
defer func() {
100105
e := recover()
101106
if e != nil {
102-
err = multierror.Append(err, xerrors.Errorf("panic in cleanup: %w", e))
107+
err = xerrors.Errorf("panic in harness.Cleanup: %+v", e)
103108
}
104109
}()
105110

106-
for _, run := range h.runs {
107-
e := run.Cleanup(ctx)
108-
if e != nil {
109-
err = multierror.Append(err, xerrors.Errorf("cleanup for %s failed: %w", run.FullID(), e))
111+
var cleanupErrs []error
112+
cleanupErrs, err = h.cleanupStrategy.Run(ctx, cleanupFns)
113+
if err != nil {
114+
err = xerrors.Errorf("cleanup strategy error: %w", err)
115+
//nolint:revive // we use named returns because we mutate it in a defer
116+
return
117+
}
118+
119+
var merr error
120+
for _, cleanupErr := range cleanupErrs {
121+
if cleanupErr != nil {
122+
merr = multierror.Append(merr, cleanupErr)
110123
}
111124
}
112125

126+
err = merr
113127
//nolint:revive // we use named returns because we mutate it in a defer
114128
return
115129
}

‎loadtest/harness/harness_test.go

+54-16
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ type panickingExecutionStrategy struct{}
1818

1919
var _ harness.ExecutionStrategy = panickingExecutionStrategy{}
2020

21-
func (panickingExecutionStrategy) Execute(_ context.Context, _ []*harness.TestRun) error {
21+
func (panickingExecutionStrategy) Run(_ context.Context, _ []harness.TestFn) ([]error, error) {
2222
panic(testPanicMessage)
2323
}
2424

@@ -28,8 +28,8 @@ type erroringExecutionStrategy struct {
2828

2929
var _ harness.ExecutionStrategy = erroringExecutionStrategy{}
3030

31-
func (e erroringExecutionStrategy) Execute(_ context.Context, _ []*harness.TestRun) error {
32-
return e.err
31+
func (e erroringExecutionStrategy) Run(_ context.Context, _ []harness.TestFn) ([]error, error) {
32+
return []error{}, e.err
3333
}
3434

3535
func Test_TestHarness(t *testing.T) {
@@ -40,7 +40,7 @@ func Test_TestHarness(t *testing.T) {
4040

4141
expectedErr := xerrors.New("expected error")
4242

43-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
43+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
4444
r1 := h.AddRun("test", "1", fakeTestFns(nil, nil))
4545
r2 := h.AddRun("test", "2", fakeTestFns(expectedErr, nil))
4646

@@ -65,7 +65,7 @@ func Test_TestHarness(t *testing.T) {
6565

6666
expectedErr := xerrors.New("expected error")
6767

68-
h := harness.NewTestHarness(erroringExecutionStrategy{err: expectedErr})
68+
h := harness.NewTestHarness(erroringExecutionStrategy{err: expectedErr}, harness.LinearExecutionStrategy{})
6969
_ = h.AddRun("test", "1", fakeTestFns(nil, nil))
7070

7171
err := h.Run(context.Background())
@@ -76,7 +76,7 @@ func Test_TestHarness(t *testing.T) {
7676
t.Run("CatchesExecutionPanic", func(t *testing.T) {
7777
t.Parallel()
7878

79-
h := harness.NewTestHarness(panickingExecutionStrategy{})
79+
h := harness.NewTestHarness(panickingExecutionStrategy{}, harness.LinearExecutionStrategy{})
8080
_ = h.AddRun("test", "1", fakeTestFns(nil, nil))
8181

8282
err := h.Run(context.Background())
@@ -93,7 +93,7 @@ func Test_TestHarness(t *testing.T) {
9393

9494
expectedErr := xerrors.New("expected error")
9595

96-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
96+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
9797
_ = h.AddRun("test", "1", fakeTestFns(nil, expectedErr))
9898

9999
err := h.Run(context.Background())
@@ -107,7 +107,7 @@ func Test_TestHarness(t *testing.T) {
107107
t.Run("Panic", func(t *testing.T) {
108108
t.Parallel()
109109

110-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
110+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
111111
_ = h.AddRun("test", "1", testFns{
112112
RunFn: func(_ context.Context, _ string, _ io.Writer) error {
113113
return nil
@@ -125,6 +125,44 @@ func Test_TestHarness(t *testing.T) {
125125
require.ErrorContains(t, err, "panic")
126126
require.ErrorContains(t, err, testPanicMessage)
127127
})
128+
129+
t.Run("CatchesExecutionError", func(t *testing.T) {
130+
t.Parallel()
131+
132+
expectedErr := xerrors.New("expected error")
133+
134+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, erroringExecutionStrategy{err: expectedErr})
135+
_ = h.AddRun("test", "1", fakeTestFns(nil, nil))
136+
137+
err := h.Run(context.Background())
138+
require.NoError(t, err)
139+
140+
err = h.Cleanup(context.Background())
141+
require.Error(t, err)
142+
require.ErrorIs(t, err, expectedErr)
143+
})
144+
145+
t.Run("CatchesExecutionPanic", func(t *testing.T) {
146+
t.Parallel()
147+
148+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, panickingExecutionStrategy{})
149+
_ = h.AddRun("test", "1", testFns{
150+
RunFn: func(_ context.Context, _ string, _ io.Writer) error {
151+
return nil
152+
},
153+
CleanupFn: func(_ context.Context, _ string) error {
154+
return nil
155+
},
156+
})
157+
158+
err := h.Run(context.Background())
159+
require.NoError(t, err)
160+
161+
err = h.Cleanup(context.Background())
162+
require.Error(t, err)
163+
require.ErrorContains(t, err, "panic")
164+
require.ErrorContains(t, err, testPanicMessage)
165+
})
128166
})
129167

130168
t.Run("Panics", func(t *testing.T) {
@@ -133,7 +171,7 @@ func Test_TestHarness(t *testing.T) {
133171
t.Run("RegisterAfterStart", func(t *testing.T) {
134172
t.Parallel()
135173

136-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
174+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
137175
_ = h.Run(context.Background())
138176

139177
require.Panics(t, func() {
@@ -144,7 +182,7 @@ func Test_TestHarness(t *testing.T) {
144182
t.Run("DuplicateTestID", func(t *testing.T) {
145183
t.Parallel()
146184

147-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
185+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
148186

149187
name, id := "test", "1"
150188
_ = h.AddRun(name, id, fakeTestFns(nil, nil))
@@ -157,7 +195,7 @@ func Test_TestHarness(t *testing.T) {
157195
t.Run("StartedTwice", func(t *testing.T) {
158196
t.Parallel()
159197

160-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
198+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
161199
h.Run(context.Background())
162200

163201
require.Panics(t, func() {
@@ -168,7 +206,7 @@ func Test_TestHarness(t *testing.T) {
168206
t.Run("ResultsBeforeStart", func(t *testing.T) {
169207
t.Parallel()
170208

171-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
209+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
172210

173211
require.Panics(t, func() {
174212
h.Results()
@@ -183,7 +221,7 @@ func Test_TestHarness(t *testing.T) {
183221
endRun = make(chan struct{})
184222
testsEnded = make(chan struct{})
185223
)
186-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
224+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
187225
_ = h.AddRun("test", "1", testFns{
188226
RunFn: func(_ context.Context, _ string, _ io.Writer) error {
189227
close(started)
@@ -210,22 +248,22 @@ func Test_TestHarness(t *testing.T) {
210248
t.Run("CleanupBeforeStart", func(t *testing.T) {
211249
t.Parallel()
212250

213-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
251+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
214252

215253
require.Panics(t, func() {
216254
h.Cleanup(context.Background())
217255
})
218256
})
219257

220-
t.Run("ClenaupBeforeFinish", func(t *testing.T) {
258+
t.Run("CleanupBeforeFinish", func(t *testing.T) {
221259
t.Parallel()
222260

223261
var (
224262
started = make(chan struct{})
225263
endRun = make(chan struct{})
226264
testsEnded = make(chan struct{})
227265
)
228-
h := harness.NewTestHarness(harness.LinearExecutionStrategy{})
266+
h := harness.NewTestHarness(harness.LinearExecutionStrategy{}, harness.LinearExecutionStrategy{})
229267
_ = h.AddRun("test", "1", testFns{
230268
RunFn: func(_ context.Context, _ string, _ io.Writer) error {
231269
close(started)

0 commit comments

Comments
 (0)