Skip to content

Commit 2e640e6

Browse files
authored
Use Resources (#5)
1 parent 5b4e3dc commit 2e640e6

28 files changed

Lines changed: 1020 additions & 109 deletions

HOOKS.md

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
# Lifecycle hooks
22

3-
testrig runs hooks at different points in the test lifecycle. **Global** hooks apply to `run`, `gotestsum`, and `diagnose`. **Iteration** hooks apply only to `diagnose`.
3+
testrig runs hooks at different points in the test lifecycle. **Global** hooks apply to the default `go test` invocation, `gotestsum`, and `diagnose`. **Iteration** hooks apply only to `diagnose`.
44

5-
Register hooks via CLI flags (before the subcommand) or Go options in a [`tools/test`](./tools/test/README.md) binary. Hooks receive `context.Context` and should respect cancellation (SIGINT cancels in-flight setup).
5+
Register hooks via CLI flags or Go options in a [`tools/test`](./tools/test/README.md) binary. Hooks receive `context.Context` and should respect cancellation (SIGINT cancels in-flight setup).
66

77
The table below is generated from [`internal/hooks/catalog.go`](./internal/hooks/catalog.go) and godoc on [`hooks.go`](./hooks.go). Run `go generate` after changes.
88

@@ -21,10 +21,11 @@ The table below is generated from [`internal/hooks/catalog.go`](./internal/hooks
2121

2222
## Flag placement
2323

24-
Put root flags **before** the subcommand:
24+
For the default `go test` invocation, testrig parses its own flags (`--ai-output`, lifecycle hooks) and forwards unknown flags to `go test`:
2525

2626
```sh
27+
testrig --global-setup "docker compose up -d" -v -count=1 ./...
2728
testrig --global-setup "docker compose up -d" diagnose --iterations 5 -- ./...
2829
```
2930

30-
`run` and `gotestsum` do not parse flags themselves; flags placed after those subcommands are forwarded to `go test`.
31+
For `gotestsum`, put root flags before the subcommand; flags after `gotestsum` are forwarded to gotestsum / `go test`.

README.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@ Run setup and teardown scripts during your test lifecycle.
4242
### CLI
4343

4444
```sh
45+
# Vanilla go test with global hooks
46+
testrig --global-setup "docker compose up -d" -v -count=1 ./...
47+
4548
# Spin up dependencies before any test, tear down after
4649
testrig diagnose --iterations 10 \
4750
--global-setup "docker compose up -d" \

go.mod

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,9 @@ require (
2121
require (
2222
github.com/bitfield/gotestdox v0.2.2 // indirect
2323
github.com/charmbracelet/colorprofile v0.4.3 // indirect
24-
github.com/charmbracelet/ultraviolet v0.0.0-20260525132238-948f4557a654 // indirect
24+
github.com/charmbracelet/ultraviolet v0.0.0-20260601155805-6cf7526a1b3f // indirect
2525
github.com/charmbracelet/x/ansi v0.11.7 // indirect
26-
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260525135217-abeec2b8bf0b // indirect
26+
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260601140041-e36ca9b00c4e // indirect
2727
github.com/charmbracelet/x/termios v0.1.1 // indirect
2828
github.com/charmbracelet/x/windows v0.2.2 // indirect
2929
github.com/clipperhouse/displaywidth v0.11.0 // indirect
@@ -38,7 +38,7 @@ require (
3838
github.com/lucasb-eyer/go-colorful v1.4.0 // indirect
3939
github.com/mattn/go-colorable v0.1.14 // indirect
4040
github.com/mattn/go-isatty v0.0.22 // indirect
41-
github.com/mattn/go-runewidth v0.0.23 // indirect
41+
github.com/mattn/go-runewidth v0.0.24 // indirect
4242
github.com/muesli/cancelreader v0.2.2 // indirect
4343
github.com/muesli/mango v0.2.0 // indirect
4444
github.com/muesli/mango-cobra v1.3.0 // indirect

go.sum

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,14 @@ github.com/charmbracelet/colorprofile v0.4.3 h1:QPa1IWkYI+AOB+fE+mg/5/4HRMZcaXex
1010
github.com/charmbracelet/colorprofile v0.4.3/go.mod h1:/zT4BhpD5aGFpqQQqw7a+VtHCzu+zrQtt1zhMt9mR4Q=
1111
github.com/charmbracelet/ultraviolet v0.0.0-20260525132238-948f4557a654 h1:FpSYhY28ucg9ZRr+2wj67FAQ0Ey5yiK0072PmRDJNek=
1212
github.com/charmbracelet/ultraviolet v0.0.0-20260525132238-948f4557a654/go.mod h1:hFpumms29Smx3LStRfku8vcCTBe1Kq8aCXtHUJa3mjY=
13+
github.com/charmbracelet/ultraviolet v0.0.0-20260601155805-6cf7526a1b3f h1:vKsPSlO4g4jKfJ9enESgNZ45BkbHngTIq3UxNOzic74=
14+
github.com/charmbracelet/ultraviolet v0.0.0-20260601155805-6cf7526a1b3f/go.mod h1:hFpumms29Smx3LStRfku8vcCTBe1Kq8aCXtHUJa3mjY=
1315
github.com/charmbracelet/x/ansi v0.11.7 h1:kzv1kJvjg2S3r9KHo8hDdHFQLEqn4RBCb39dAYC84jI=
1416
github.com/charmbracelet/x/ansi v0.11.7/go.mod h1:9qGpnAVYz+8ACONkZBUWPtL7lulP9No6p1epAihUZwQ=
1517
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260525135217-abeec2b8bf0b h1:8eaOsqj4avmPlTVXAQE9Z/SWqKc5imsBXdcdkaZb1QE=
1618
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260525135217-abeec2b8bf0b/go.mod h1:nsExn0DGyX0lh9LwLHTn2Gg+hafdzfSXnC+QmEJTZFY=
19+
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260601140041-e36ca9b00c4e h1:FuUQQGe5VgBkk0ssy7O1CB9x366HGrL1LycUeZRYl10=
20+
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260601140041-e36ca9b00c4e/go.mod h1:nsExn0DGyX0lh9LwLHTn2Gg+hafdzfSXnC+QmEJTZFY=
1721
github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f h1:pk6gmGpCE7F3FcjaOEKYriCvpmIN4+6OS/RD0vm4uIA=
1822
github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f/go.mod h1:IfZAMTHB6XkZSeXUqriemErjAWCCzT0LwjKFYCZyw0I=
1923
github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
@@ -57,6 +61,8 @@ github.com/mattn/go-isatty v0.0.22 h1:j8l17JJ9i6VGPUFUYoTUKPSgKe/83EYU2zBC7YNKMw
5761
github.com/mattn/go-isatty v0.0.22/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4=
5862
github.com/mattn/go-runewidth v0.0.23 h1:7ykA0T0jkPpzSvMS5i9uoNn2Xy3R383f9HDx3RybWcw=
5963
github.com/mattn/go-runewidth v0.0.23/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
64+
github.com/mattn/go-runewidth v0.0.24 h1:cpokDiIn0MGnhdHwuWnJBITySJ20QyNGnY2kR/ay2DU=
65+
github.com/mattn/go-runewidth v0.0.24/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
6066
github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
6167
github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
6268
github.com/muesli/mango v0.2.0 h1:iNNc0c5VLQ6fsMgAqGQofByNUBH2Q2nEbD6TaI+5yyQ=

hooks.go

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@
55
package testrig
66

77
import (
8+
"github.com/spf13/cobra"
9+
"github.com/spf13/pflag"
10+
811
"github.com/smartcontractkit/testrig/internal/cmd"
912
"github.com/smartcontractkit/testrig/internal/hooks"
1013
)
@@ -20,6 +23,13 @@ type Option = hooks.Option
2023
// It is exported for internal use by the CLI engine.
2124
type RunOptions = hooks.RunOptions
2225

26+
// Resource is one prepared, isolated piece of infrastructure (e.g. a database)
27+
// supplied by a ResourceProvider. See hooks.Resource for field semantics.
28+
type Resource = hooks.Resource
29+
30+
// ResourceProvider supplies isolated resources for a run. See WithResources.
31+
type ResourceProvider = hooks.ResourceProvider
32+
2333
// GlobalSetup registers a hook to run once before any tests.
2434
func GlobalSetup(h Hook) Option {
2535
return hooks.GlobalSetup(h)
@@ -48,6 +58,34 @@ func IterationTeardown(h Hook) Option {
4858
return hooks.IterationTeardown(h)
4959
}
5060

61+
// WithResources registers a provider that supplies isolated infrastructure
62+
// (e.g. databases) for the run. See ResourceProvider for count semantics and
63+
// partial-failure behavior. Each resource's Env is appended to the child go test
64+
// process; Reset and DumpDiagnostics apply during diagnose; Cleanup runs once
65+
// after the command finishes.
66+
func WithResources(p ResourceProvider) Option {
67+
return hooks.WithResources(p)
68+
}
69+
70+
// WithCommand registers an additional subcommand on the testrig root command,
71+
// for project-specific utilities (e.g. persistent database management).
72+
func WithCommand(cmd *cobra.Command) Option {
73+
return hooks.WithCommand(cmd)
74+
}
75+
76+
// WithRootFlags registers persistent flags on the root command, available to
77+
// every subcommand. Use it to add consumer flags (e.g. --database-url) that a
78+
// resource provider or custom command reads.
79+
func WithRootFlags(register func(*pflag.FlagSet)) Option {
80+
return hooks.WithRootFlags(register)
81+
}
82+
83+
// WithRootCommand sets the CLI name used in help text, examples, and
84+
// cobra.CommandPath(). Defaults to "testrig" when unset.
85+
func WithRootCommand(name string) Option {
86+
return hooks.WithRootCommand(name)
87+
}
88+
5189
// BuildOptions evaluates the functional options and returns the internal struct.
5290
// It is exported for internal use by the CLI engine.
5391
func BuildOptions(opts ...Option) RunOptions {

internal/cmd/args.go

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
package cmd
2+
3+
import (
4+
"strings"
5+
6+
"github.com/spf13/cobra"
7+
"github.com/spf13/pflag"
8+
)
9+
10+
// isHelpRequest reports whether args ask for testrig usage (not go test). Root uses
11+
// DisableFlagParsing, so cobra does not intercept -h/--help; we handle that before
12+
// forwarding to go test. Args after "--" are left for go test (e.g. testrig -- -h).
13+
func isHelpRequest(args []string) bool {
14+
for _, arg := range args {
15+
if arg == "--" {
16+
return false
17+
}
18+
switch arg {
19+
case "-h", "--help":
20+
return true
21+
}
22+
}
23+
return false
24+
}
25+
26+
// goTestArgsFromRoot parses testrig persistent flags from args, applies them to the
27+
// root flag set, and returns the remainder for go test. Root uses DisableFlagParsing
28+
// so go test flags (-v, -count, -run, etc.) are never interpreted by pflag.
29+
func goTestArgsFromRoot(cmd *cobra.Command, args []string) ([]string, error) {
30+
flags := cmd.Root().PersistentFlags()
31+
var testrigArgs, goTestArgs []string
32+
33+
for i := 0; i < len(args); i++ {
34+
arg := args[i]
35+
if arg == "--" {
36+
goTestArgs = append(goTestArgs, args[i:]...)
37+
break
38+
}
39+
40+
if !strings.HasPrefix(arg, "-") {
41+
goTestArgs = append(goTestArgs, arg)
42+
continue
43+
}
44+
45+
name := strings.TrimLeft(arg, "-")
46+
if before, _, found := strings.Cut(name, "="); found {
47+
name = before
48+
}
49+
50+
// Only match long flags for now to match current behavior,
51+
// but using native pflag lookup.
52+
var f *pflag.Flag
53+
if strings.HasPrefix(arg, "--") {
54+
f = flags.Lookup(name)
55+
} else if len(name) == 1 {
56+
f = flags.ShorthandLookup(name)
57+
}
58+
59+
if f != nil {
60+
switch {
61+
case strings.Contains(arg, "="):
62+
testrigArgs = append(testrigArgs, arg)
63+
case f.NoOptDefVal != "":
64+
testrigArgs = append(testrigArgs, arg)
65+
case i+1 < len(args):
66+
testrigArgs = append(testrigArgs, arg, args[i+1])
67+
i++
68+
default:
69+
testrigArgs = append(testrigArgs, arg)
70+
}
71+
continue
72+
}
73+
74+
goTestArgs = append(goTestArgs, arg)
75+
}
76+
77+
if len(testrigArgs) > 0 {
78+
if err := flags.Parse(testrigArgs); err != nil {
79+
return nil, err
80+
}
81+
}
82+
83+
return goTestArgs, nil
84+
}
85+
86+
// runRootAfterParsing applies persistent root flags from args, then runs fn inside
87+
// withGlobalHooks so global setup/teardown see the parsed flag values. Help requests
88+
// return before any hooks run.
89+
func runRootAfterParsing(cmd *cobra.Command, args []string, fn func(remainder []string) error) error {
90+
if isHelpRequest(args) {
91+
return pflag.ErrHelp
92+
}
93+
remainder, err := goTestArgsFromRoot(cmd, args)
94+
if err != nil {
95+
return err
96+
}
97+
return withGlobalHooks(func(*cobra.Command, []string) error {
98+
return fn(remainder)
99+
})(cmd, args)
100+
}

internal/cmd/args_test.go

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
package cmd
2+
3+
import (
4+
"context"
5+
"os"
6+
"path/filepath"
7+
"testing"
8+
9+
"github.com/spf13/cobra"
10+
"github.com/spf13/pflag"
11+
"github.com/stretchr/testify/assert"
12+
"github.com/stretchr/testify/require"
13+
14+
"github.com/smartcontractkit/testrig/internal/hooks"
15+
)
16+
17+
func TestIsHelpRequest(t *testing.T) {
18+
t.Parallel()
19+
20+
tests := []struct {
21+
name string
22+
args []string
23+
want bool
24+
}{
25+
{name: "short help", args: []string{"-h"}, want: true},
26+
{name: "long help", args: []string{"--help"}, want: true},
27+
{name: "go test args", args: []string{"-v", "./..."}, want: false},
28+
{name: "after separator", args: []string{"--", "-h", "./..."}, want: false},
29+
}
30+
31+
for _, tc := range tests {
32+
t.Run(tc.name, func(t *testing.T) {
33+
t.Parallel()
34+
assert.Equal(t, tc.want, isHelpRequest(tc.args))
35+
})
36+
}
37+
}
38+
39+
func newRootArgsTestCmd(t *testing.T) *cobra.Command {
40+
t.Helper()
41+
cmd := &cobra.Command{Use: "testrig"}
42+
cmd.PersistentFlags().Bool("ai-output", false, "")
43+
hooks.RegisterPersistentFlags(cmd.PersistentFlags())
44+
return cmd
45+
}
46+
47+
func TestGoTestArgsFromRoot(t *testing.T) {
48+
t.Parallel()
49+
50+
tests := []struct {
51+
name string
52+
args []string
53+
want []string
54+
wantGlobal string
55+
wantAIOut bool
56+
}{
57+
{
58+
name: "go test verbose after persistent",
59+
args: []string{"--ai-output", "-v", "./pkg"},
60+
want: []string{"-v", "./pkg"},
61+
wantAIOut: true,
62+
},
63+
{
64+
name: "go test count flag",
65+
args: []string{"-count=1", "./..."},
66+
want: []string{"-count=1", "./..."},
67+
},
68+
{
69+
name: "go test run flag",
70+
args: []string{"-run", "^TestFoo$", "./pkg"},
71+
want: []string{"-run", "^TestFoo$", "./pkg"},
72+
},
73+
{
74+
name: "global setup then go test flags",
75+
args: []string{"--global-setup", "docker compose up -d", "-run", "^TestFoo$", "./pkg"},
76+
want: []string{"-run", "^TestFoo$", "./pkg"},
77+
wantGlobal: "docker compose up -d",
78+
},
79+
{
80+
name: "global setup equals form",
81+
args: []string{"--global-setup=docker compose up -d", "-v", "./..."},
82+
want: []string{"-v", "./..."},
83+
wantGlobal: "docker compose up -d",
84+
},
85+
{
86+
name: "double dash forwards remainder including go test help",
87+
args: []string{"--ai-output", "--", "-h", "./pkg"},
88+
want: []string{"--", "-h", "./pkg"},
89+
wantAIOut: true,
90+
},
91+
}
92+
93+
for _, tc := range tests {
94+
t.Run(tc.name, func(t *testing.T) {
95+
t.Parallel()
96+
cmd := newRootArgsTestCmd(t)
97+
98+
got, err := goTestArgsFromRoot(cmd, tc.args)
99+
require.NoError(t, err)
100+
assert.Equal(t, tc.want, got)
101+
102+
if tc.wantAIOut {
103+
v, err := cmd.PersistentFlags().GetBool("ai-output")
104+
require.NoError(t, err)
105+
assert.True(t, v)
106+
}
107+
if tc.wantGlobal != "" {
108+
v, err := cmd.PersistentFlags().GetString("global-setup")
109+
require.NoError(t, err)
110+
assert.Equal(t, tc.wantGlobal, v)
111+
}
112+
})
113+
}
114+
}
115+
116+
func TestRunRootAfterParsingAppliesFlagsBeforeHooks(t *testing.T) {
117+
t.Parallel()
118+
119+
dir := t.TempDir()
120+
marker := filepath.Join(dir, "setup")
121+
cmd := newRootArgsTestCmd(t)
122+
cmd.SetContext(context.Background())
123+
124+
err := runRootAfterParsing(cmd, []string{"--global-setup", touchCmd(marker)}, func([]string) error {
125+
return nil
126+
})
127+
require.NoError(t, err)
128+
assert.FileExists(t, marker)
129+
}
130+
131+
func TestRunRootAfterParsingHelpSkipsHooks(t *testing.T) {
132+
t.Parallel()
133+
134+
dir := t.TempDir()
135+
marker := filepath.Join(dir, "setup")
136+
cmd := newRootArgsTestCmd(t)
137+
cmd.SetContext(context.Background())
138+
139+
err := runRootAfterParsing(cmd, []string{"--global-setup", touchCmd(marker), "-h"}, func([]string) error {
140+
return nil
141+
})
142+
require.ErrorIs(t, err, pflag.ErrHelp)
143+
_, statErr := os.Stat(marker)
144+
assert.True(t, os.IsNotExist(statErr))
145+
}

0 commit comments

Comments
 (0)