From 39fc242a3c0e3d02349d69c791e9f084f7e536ab Mon Sep 17 00:00:00 2001
From: Steven Hartland <stevenmhartland@gmail.com>
Date: Wed, 21 Feb 2024 22:54:01 +0000
Subject: [PATCH] feat: WithEnv customize request option

Add WithEnv which set additional environment variables on a request
so that callers can customise a request.

This can be used to easily add environment variables to a container
created using one of the standard module e.g. postgres.
---
 docs/features/common_functional_options.md | 10 +++++
 options.go                                 | 14 +++++++
 options_test.go                            | 48 ++++++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/docs/features/common_functional_options.md b/docs/features/common_functional_options.md
index e9f35bcaff4..6a6e4120244 100644
--- a/docs/features/common_functional_options.md
+++ b/docs/features/common_functional_options.md
@@ -15,6 +15,16 @@ _Testcontainers for Go_ exposes an interface to perform this operations: `ImageS
 
 Using the `WithImageSubstitutors` options, you could define your own substitutions to the container images. E.g. adding a prefix to the images so that they can be pulled from a Docker registry other than Docker Hub. This is the usual mechanism for using Docker image proxies, caches, etc.
 
+#### WithEnv
+
+- Since testcontainers-go <a href="https://github.com/testcontainers/testcontainers-go/releases/tag/v0.29.0"><span class="tc-version">:material-tag: v0.29.0</span></a>
+
+If you need to either pass additional environment variables to a container or override them, you can use `testcontainers.WithEnv` for example:
+
+```golang
+postgres, err = postgresModule.RunContainer(ctx, testcontainers.WithEnv(map[string]string{"POSTGRES_INITDB_ARGS", "--no-sync"}))
+```
+
 #### WithLogConsumers
 
 - Since testcontainers-go <a href="https://github.com/testcontainers/testcontainers-go/releases/tag/v0.28.0"><span class="tc-version">:material-tag: v0.28.0</span></a>
diff --git a/options.go b/options.go
index 3811e4eb536..34881dba85a 100644
--- a/options.go
+++ b/options.go
@@ -53,6 +53,20 @@ func WithEndpointSettingsModifier(modifier func(settings map[string]*network.End
 	}
 }
 
+// WithEnv sets the environment variables for a container.
+// If the environment variable already exists, it will be overridden.
+func WithEnv(envs map[string]string) CustomizeRequestOption {
+	return func(req *GenericContainerRequest) {
+		if req.Env == nil {
+			req.Env = map[string]string{}
+		}
+
+		for key, val := range envs {
+			req.Env[key] = val
+		}
+	}
+}
+
 // WithHostConfigModifier allows to override the default host config
 func WithHostConfigModifier(modifier func(hostConfig *container.HostConfig)) CustomizeRequestOption {
 	return func(req *GenericContainerRequest) {
diff --git a/options_test.go b/options_test.go
index e13642ac099..bbd4a944d5c 100644
--- a/options_test.go
+++ b/options_test.go
@@ -162,3 +162,51 @@ func TestWithAfterReadyCommand(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "/tmp/.testcontainers\n", string(content))
 }
+
+func TestWithEnv(t *testing.T) {
+	tests := map[string]struct {
+		req    *testcontainers.GenericContainerRequest
+		env    map[string]string
+		expect map[string]string
+	}{
+		"add": {
+			req: &testcontainers.GenericContainerRequest{
+				ContainerRequest: testcontainers.ContainerRequest{
+					Env: map[string]string{"KEY1": "VAL1"},
+				},
+			},
+			env: map[string]string{"KEY2": "VAL2"},
+			expect: map[string]string{
+				"KEY1": "VAL1",
+				"KEY2": "VAL2",
+			},
+		},
+		"add-nil": {
+			req:    &testcontainers.GenericContainerRequest{},
+			env:    map[string]string{"KEY2": "VAL2"},
+			expect: map[string]string{"KEY2": "VAL2"},
+		},
+		"override": {
+			req: &testcontainers.GenericContainerRequest{
+				ContainerRequest: testcontainers.ContainerRequest{
+					Env: map[string]string{
+						"KEY1": "VAL1",
+						"KEY2": "VAL2",
+					},
+				},
+			},
+			env: map[string]string{"KEY2": "VAL3"},
+			expect: map[string]string{
+				"KEY1": "VAL1",
+				"KEY2": "VAL3",
+			},
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			opt := testcontainers.WithEnv(tc.env)
+			opt.Customize(tc.req)
+			require.Equal(t, tc.expect, tc.req.Env)
+		})
+	}
+}