Skip to content

Commit 058b08d

Browse files
mdelapenyaclaude
andauthored
chore(k6|localstack|kafka|mariadb): use Run function (#3414)
* chore(k3s): use WithFiles * chore(k6): use Run function 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * chore(localstack): use Run function 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * chore(localstack): use Run function 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * chore(kafka): use Run function 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * chore(mariadb): use Run function 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * chore(kafka): more accurate CMD * chore(kafka): more accurate initial lifecycle hooks * fix(kafka): avoid panics in version validation --------- Co-authored-by: Claude <[email protected]>
1 parent d68b675 commit 058b08d

File tree

9 files changed

+223
-233
lines changed

9 files changed

+223
-233
lines changed

modules/k3s/k3s.go

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,10 @@ func WithManifest(manifestPath string) testcontainers.CustomizeRequestOption {
3939
manifest := filepath.Base(manifestPath)
4040
target := k3sManifests + manifest
4141

42-
req.Files = append(req.Files, testcontainers.ContainerFile{
42+
return testcontainers.WithFiles(testcontainers.ContainerFile{
4343
HostFilePath: manifestPath,
4444
ContainerFilePath: target,
45-
})
46-
47-
return nil
45+
})(req)
4846
}
4947
}
5048

modules/k6/k6.go

Lines changed: 25 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -86,19 +86,17 @@ func WithTestScript(scriptPath string) testcontainers.CustomizeRequestOption {
8686
func WithTestScriptReader(reader io.Reader, scriptBaseName string) testcontainers.CustomizeRequestOption {
8787
opt := func(req *testcontainers.GenericContainerRequest) error {
8888
target := "/home/k6x/" + scriptBaseName
89-
req.Files = append(
90-
req.Files,
91-
testcontainers.ContainerFile{
92-
Reader: reader,
93-
ContainerFilePath: target,
94-
FileMode: 0o644,
95-
},
96-
)
9789

98-
// add script to the k6 run command
99-
req.Cmd = append(req.Cmd, target)
90+
if err := testcontainers.WithFiles(testcontainers.ContainerFile{
91+
Reader: reader,
92+
ContainerFilePath: target,
93+
FileMode: 0o644,
94+
})(req); err != nil {
95+
return err
96+
}
10097

101-
return nil
98+
// add script to the k6 run command
99+
return testcontainers.WithCmdArgs(target)(req)
102100
}
103101
return opt
104102
}
@@ -123,11 +121,7 @@ func WithCmdOptions(options ...string) testcontainers.CustomizeRequestOption {
123121

124122
// SetEnvVar adds a '--env' command-line flag to the k6 command in the container for setting an environment variable for the test script.
125123
func SetEnvVar(variable string, value string) testcontainers.CustomizeRequestOption {
126-
return func(req *testcontainers.GenericContainerRequest) error {
127-
req.Cmd = append(req.Cmd, "--env", fmt.Sprintf("%s=%s", variable, value))
128-
129-
return nil
130-
}
124+
return testcontainers.WithCmdArgs("--env", variable+"="+value)
131125
}
132126

133127
// WithCache sets a volume as a cache for building the k6 binary
@@ -146,18 +140,13 @@ func WithCache() testcontainers.CustomizeRequestOption {
146140
}
147141
}
148142

149-
return func(req *testcontainers.GenericContainerRequest) error {
150-
mount := testcontainers.ContainerMount{
151-
Source: testcontainers.DockerVolumeMountSource{
152-
Name: cacheVol,
153-
VolumeOptions: volOptions,
154-
},
155-
Target: cacheTarget,
156-
}
157-
req.Mounts = append(req.Mounts, mount)
158-
159-
return nil
160-
}
143+
return testcontainers.WithMounts(testcontainers.ContainerMount{
144+
Source: testcontainers.DockerVolumeMountSource{
145+
Name: cacheVol,
146+
VolumeOptions: volOptions,
147+
},
148+
Target: cacheTarget,
149+
})
161150
}
162151

163152
// Deprecated: use Run instead
@@ -168,31 +157,20 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize
168157

169158
// Run creates an instance of the K6 container type
170159
func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustomizer) (*K6Container, error) {
171-
req := testcontainers.ContainerRequest{
172-
Image: img,
173-
Cmd: []string{"run"},
174-
WaitingFor: wait.ForExit(),
175-
}
176-
177-
genericContainerReq := testcontainers.GenericContainerRequest{
178-
ContainerRequest: req,
179-
Started: true,
160+
moduleOpts := []testcontainers.ContainerCustomizer{
161+
testcontainers.WithCmdArgs("run"),
162+
testcontainers.WithWaitStrategy(wait.ForExit()),
180163
}
181164

182-
for _, opt := range opts {
183-
if err := opt.Customize(&genericContainerReq); err != nil {
184-
return nil, err
185-
}
186-
}
165+
moduleOpts = append(moduleOpts, opts...)
187166

188-
container, err := testcontainers.GenericContainer(ctx, genericContainerReq)
189167
var c *K6Container
190-
if container != nil {
191-
c = &K6Container{Container: container}
168+
ctr, err := testcontainers.Run(ctx, img, moduleOpts...)
169+
if ctr != nil {
170+
c = &K6Container{Container: ctr}
192171
}
193-
194172
if err != nil {
195-
return c, fmt.Errorf("generic container: %w", err)
173+
return c, fmt.Errorf("run k6: %w", err)
196174
}
197175

198176
return c, nil

modules/kafka/kafka.go

Lines changed: 70 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,13 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize
4646

4747
// Run creates an instance of the Kafka container type
4848
func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustomizer) (*KafkaContainer, error) {
49-
req := testcontainers.ContainerRequest{
50-
Image: img,
51-
ExposedPorts: []string{string(publicPort)},
52-
Env: map[string]string{
49+
if err := validateKRaftVersion(img); err != nil {
50+
return nil, err
51+
}
52+
53+
moduleOpts := []testcontainers.ContainerCustomizer{
54+
testcontainers.WithExposedPorts(string(publicPort)),
55+
testcontainers.WithEnv(map[string]string{
5356
// envVars {
5457
"KAFKA_LISTENERS": "PLAINTEXT://0.0.0.0:9093,BROKER://0.0.0.0:9092,CONTROLLER://0.0.0.0:9094",
5558
"KAFKA_REST_BOOTSTRAP_SERVERS": "PLAINTEXT://0.0.0.0:9093,BROKER://0.0.0.0:9092,CONTROLLER://0.0.0.0:9094",
@@ -66,56 +69,53 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom
6669
"KAFKA_PROCESS_ROLES": "broker,controller",
6770
"KAFKA_CONTROLLER_LISTENER_NAMES": "CONTROLLER",
6871
// }
69-
},
70-
Entrypoint: []string{"sh"},
72+
}),
73+
testcontainers.WithEntrypoint("sh"),
7174
// this CMD will wait for the starter script to be copied into the container and then execute it
72-
Cmd: []string{"-c", "while [ ! -f " + starterScript + " ]; do sleep 0.1; done; bash " + starterScript},
73-
LifecycleHooks: []testcontainers.ContainerLifecycleHooks{
74-
{
75-
PostStarts: []testcontainers.ContainerHook{
76-
// Use a single hook to copy the starter script and wait for
77-
// the Kafka server to be ready. This prevents the wait running
78-
// if the starter script fails to copy.
79-
func(ctx context.Context, c testcontainers.Container) error {
80-
// 1. copy the starter script into the container
81-
if err := copyStarterScript(ctx, c); err != nil {
82-
return fmt.Errorf("copy starter script: %w", err)
83-
}
84-
85-
// 2. wait for the Kafka server to be ready
86-
return wait.ForLog(".*Transitioning from RECOVERY to RUNNING.*").AsRegexp().WaitUntilReady(ctx, c)
87-
},
75+
testcontainers.WithCmd("-c", "while [ ! -f "+starterScript+" ]; do sleep 0.1; done; bash "+starterScript),
76+
testcontainers.WithLifecycleHooks(testcontainers.ContainerLifecycleHooks{
77+
PostStarts: []testcontainers.ContainerHook{
78+
// Use a single hook to copy the starter script and wait for
79+
// the Kafka server to be ready. This prevents the wait running
80+
// if the starter script fails to copy.
81+
func(ctx context.Context, c testcontainers.Container) error {
82+
// 1. copy the starter script into the container
83+
if err := copyStarterScript(ctx, c); err != nil {
84+
return fmt.Errorf("copy starter script: %w", err)
85+
}
86+
87+
// 2. wait for the Kafka server to be ready
88+
return wait.ForLog(".*Transitioning from RECOVERY to RUNNING.*").AsRegexp().WaitUntilReady(ctx, c)
8889
},
8990
},
90-
},
91+
}),
9192
}
9293

93-
genericContainerReq := testcontainers.GenericContainerRequest{
94-
ContainerRequest: req,
95-
Started: true,
96-
}
94+
moduleOpts = append(moduleOpts, opts...)
9795

98-
for _, opt := range opts {
99-
if err := opt.Customize(&genericContainerReq); err != nil {
100-
return nil, err
101-
}
102-
}
96+
// configure the controller quorum voters after all the options have been applied
97+
moduleOpts = append(moduleOpts, configureControllerQuorumVoters())
10398

104-
err := validateKRaftVersion(genericContainerReq.Image)
99+
var c *KafkaContainer
100+
ctr, err := testcontainers.Run(ctx, img, moduleOpts...)
101+
if ctr != nil {
102+
c = &KafkaContainer{Container: ctr}
103+
}
105104
if err != nil {
106-
return nil, err
105+
return c, fmt.Errorf("run kafka: %w", err)
107106
}
108107

109-
configureControllerQuorumVoters(&genericContainerReq)
110-
111-
container, err := testcontainers.GenericContainer(ctx, genericContainerReq)
112-
var c *KafkaContainer
113-
if container != nil {
114-
c = &KafkaContainer{Container: container, ClusterID: genericContainerReq.Env["CLUSTER_ID"]}
108+
// Inspect the container to get the CLUSTER_ID environment variable
109+
inspect, err := ctr.Inspect(ctx)
110+
if err != nil {
111+
return c, fmt.Errorf("inspect kafka: %w", err)
115112
}
116113

117-
if err != nil {
118-
return c, fmt.Errorf("generic container: %w", err)
114+
for _, env := range inspect.Config.Env {
115+
if v, ok := strings.CutPrefix(env, "CLUSTER_ID="); ok {
116+
c.ClusterID = v
117+
break
118+
}
119119
}
120120

121121
return c, nil
@@ -150,11 +150,9 @@ func copyStarterScript(ctx context.Context, c testcontainers.Container) error {
150150
}
151151

152152
func WithClusterID(clusterID string) testcontainers.CustomizeRequestOption {
153-
return func(req *testcontainers.GenericContainerRequest) error {
154-
req.Env["CLUSTER_ID"] = clusterID
155-
156-
return nil
157-
}
153+
return testcontainers.WithEnv(map[string]string{
154+
"CLUSTER_ID": clusterID,
155+
})
158156
}
159157

160158
// Brokers retrieves the broker connection strings from Kafka with only one entry,
@@ -168,24 +166,28 @@ func (kc *KafkaContainer) Brokers(ctx context.Context) ([]string, error) {
168166
return []string{endpoint}, nil
169167
}
170168

171-
// configureControllerQuorumVoters sets the quorum voters for the controller. For that, it will
172-
// check if there are any network aliases defined for the container and use the first alias in the
173-
// first network. Else, it will use localhost.
174-
func configureControllerQuorumVoters(req *testcontainers.GenericContainerRequest) {
175-
if req.Env == nil {
176-
req.Env = map[string]string{}
177-
}
169+
// configureControllerQuorumVoters returns an option that sets the quorum voters for the controller.
170+
// For that, it will check if there are any network aliases defined for the container and use the
171+
// first alias in the first network. Else, it will use localhost.
172+
func configureControllerQuorumVoters() testcontainers.CustomizeRequestOption {
173+
return func(req *testcontainers.GenericContainerRequest) error {
174+
if req.Env == nil {
175+
req.Env = map[string]string{}
176+
}
178177

179-
if req.Env["KAFKA_CONTROLLER_QUORUM_VOTERS"] == "" {
180-
host := "localhost"
181-
if len(req.Networks) > 0 {
182-
nw := req.Networks[0]
183-
if len(req.NetworkAliases[nw]) > 0 {
184-
host = req.NetworkAliases[nw][0]
178+
if req.Env["KAFKA_CONTROLLER_QUORUM_VOTERS"] == "" {
179+
host := "localhost"
180+
if len(req.Networks) > 0 {
181+
nw := req.Networks[0]
182+
if len(req.NetworkAliases[nw]) > 0 {
183+
host = req.NetworkAliases[nw][0]
184+
}
185185
}
186+
187+
req.Env["KAFKA_CONTROLLER_QUORUM_VOTERS"] = "1@" + host + ":9094"
186188
}
187189

188-
req.Env["KAFKA_CONTROLLER_QUORUM_VOTERS"] = fmt.Sprintf("1@%s:9094", host)
190+
return nil
189191
}
190192
// }
191193
}
@@ -197,8 +199,13 @@ func validateKRaftVersion(fqName string) error {
197199
return errors.New("image cannot be empty")
198200
}
199201

200-
image := fqName[:strings.LastIndex(fqName, ":")]
201-
version := fqName[strings.LastIndex(fqName, ":")+1:]
202+
idx := strings.LastIndex(fqName, ":")
203+
if idx == -1 || idx == len(fqName)-1 {
204+
return nil
205+
}
206+
207+
image := fqName[:idx]
208+
version := fqName[idx+1:]
202209

203210
if !strings.EqualFold(image, "confluentinc/confluent-local") {
204211
// do not validate if the image is not the official one.

modules/kafka/kafka_helpers_test.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,8 @@ func TestConfigureQuorumVoters(t *testing.T) {
5555

5656
for _, test := range tests {
5757
t.Run(test.name, func(t *testing.T) {
58-
configureControllerQuorumVoters(test.req)
58+
err := configureControllerQuorumVoters()(test.req)
59+
require.NoError(t, err)
5960

6061
require.Equalf(t, test.expectedVoters, test.req.Env["KAFKA_CONTROLLER_QUORUM_VOTERS"], "expected KAFKA_CONTROLLER_QUORUM_VOTERS to be %s, got %s", test.expectedVoters, test.req.Env["KAFKA_CONTROLLER_QUORUM_VOTERS"])
6162
})
@@ -93,6 +94,11 @@ func TestValidateKRaftVersion(t *testing.T) {
9394
image: "my-kafka:1.0.0",
9495
wantErr: false,
9596
},
97+
{
98+
name: "lacks tag",
99+
image: "my-kafka",
100+
wantErr: false,
101+
},
96102
}
97103

98104
for _, test := range tests {

0 commit comments

Comments
 (0)