Skip to content

Commit 0deabb1

Browse files
committed
feat: config to override env vars
Add a config option to env provider to allow users to change the source of environment variables. Doing so is particularly useful in parallel tests where t.Setenv will not work.
1 parent 4a21f9b commit 0deabb1

File tree

2 files changed

+159
-28
lines changed

2 files changed

+159
-28
lines changed

providers/env/env.go

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,18 @@ type Env struct {
1515
prefix string
1616
delim string
1717
cb func(key string, value string) (string, interface{})
18+
opt *Opt
19+
}
20+
21+
// Opt represents optional configuration passed to the provider.
22+
type Opt struct {
23+
// EnvironFunc is the function that feeds environment variables
24+
// to the provider.
25+
EnvironFunc func() []string
26+
}
27+
28+
var defaultOpt = &Opt{
29+
EnvironFunc: os.Environ,
1830
}
1931

2032
// Provider returns an environment variables provider that returns
@@ -30,29 +42,48 @@ type Env struct {
3042
// everything, strip prefixes and replace _ with . etc.
3143
// If the callback returns an empty string, the variable will be
3244
// ignored.
33-
func Provider(prefix, delim string, cb func(s string) string) *Env {
45+
//
46+
// It takes an optional Opt argument containing a function to override
47+
// the default source for environment variables, which can be useful
48+
// for mocking and parallel unit tests.
49+
func Provider(prefix, delim string, cb func(s string) string, opt ...*Opt) *Env {
3450
e := &Env{
3551
prefix: prefix,
3652
delim: delim,
53+
opt: defaultOpt,
3754
}
3855
if cb != nil {
3956
e.cb = func(key string, value string) (string, interface{}) {
4057
return cb(key), value
4158
}
4259
}
60+
if len(opt) > 0 {
61+
e.opt = opt[0]
62+
}
63+
4364
return e
4465
}
4566

4667
// ProviderWithValue works exactly the same as Provider except the callback
4768
// takes a (key, value) with the variable name and value and allows you
4869
// to modify both. This is useful for cases where you may want to return
4970
// other types like a string slice instead of just a string.
50-
func ProviderWithValue(prefix, delim string, cb func(key string, value string) (string, interface{})) *Env {
51-
return &Env{
71+
//
72+
// It takes an optional Opt argument containing a function to override
73+
// the default source for environment variables, which can be useful
74+
// for mocking and parallel unit tests.
75+
func ProviderWithValue(prefix, delim string, cb func(key string, value string) (string, interface{}), opt ...*Opt) *Env {
76+
e := &Env{
5277
prefix: prefix,
5378
delim: delim,
5479
cb: cb,
80+
opt: defaultOpt,
5581
}
82+
if len(opt) > 0 {
83+
e.opt = opt[0]
84+
}
85+
86+
return e
5687
}
5788

5889
// ReadBytes is not supported by the env provider.
@@ -65,7 +96,7 @@ func (e *Env) ReadBytes() ([]byte, error) {
6596
func (e *Env) Read() (map[string]interface{}, error) {
6697
// Collect the environment variable keys.
6798
var keys []string
68-
for _, k := range os.Environ() {
99+
for _, k := range e.opt.EnvironFunc() {
69100
if e.prefix != "" {
70101
if strings.HasPrefix(k, e.prefix) {
71102
keys = append(keys, k)

providers/env/env_test.go

Lines changed: 124 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
package env
22

33
import (
4-
"github.com/stretchr/testify/assert"
5-
"os"
4+
"slices"
65
"strings"
76
"testing"
7+
8+
"github.com/stretchr/testify/assert"
89
)
910

1011
func TestProvider(t *testing.T) {
12+
mockEnviron := func() []string {
13+
return []string{"TEST_FOO=bar"}
14+
}
1115

1216
testCases := []struct {
1317
name string
@@ -18,6 +22,7 @@ func TestProvider(t *testing.T) {
1822
expKey string
1923
expValue string
2024
cb func(key string) string
25+
opt *Opt
2126
want *Env
2227
}{
2328
{
@@ -27,6 +32,7 @@ func TestProvider(t *testing.T) {
2732
want: &Env{
2833
prefix: "TESTVAR_",
2934
delim: ".",
35+
opt: defaultOpt,
3036
},
3137
},
3238
{
@@ -43,6 +49,7 @@ func TestProvider(t *testing.T) {
4349
want: &Env{
4450
prefix: "TESTVAR_",
4551
delim: ".",
52+
opt: defaultOpt,
4653
},
4754
},
4855
{
@@ -52,6 +59,7 @@ func TestProvider(t *testing.T) {
5259
want: &Env{
5360
prefix: "",
5461
delim: ".",
62+
opt: defaultOpt,
5563
},
5664
},
5765
{
@@ -66,50 +74,87 @@ func TestProvider(t *testing.T) {
6674
return strings.Replace(strings.ToUpper(key), "_", ".", -1)
6775
},
6876
},
77+
{
78+
name: "Custom opt",
79+
prefix: "TEST_",
80+
delim: ".",
81+
opt: &Opt{
82+
EnvironFunc: mockEnviron,
83+
},
84+
want: &Env{
85+
prefix: "TEST_",
86+
delim: ".",
87+
opt: &Opt{
88+
EnvironFunc: mockEnviron,
89+
},
90+
},
91+
},
6992
}
7093

7194
for _, tc := range testCases {
7295
t.Run(tc.name, func(t *testing.T) {
73-
gotProvider := Provider(tc.prefix, tc.delim, tc.cb)
74-
if tc.cb == nil {
96+
var gotProvider *Env
97+
if tc.opt != nil {
98+
gotProvider = Provider(tc.prefix, tc.delim, tc.cb, tc.opt)
99+
} else {
100+
gotProvider = Provider(tc.prefix, tc.delim, tc.cb)
101+
}
102+
103+
if tc.cb == nil && tc.opt == nil {
75104
assert.Equal(t, tc.want, gotProvider)
105+
return
76106
}
107+
77108
if tc.cb != nil {
78109
k, v := gotProvider.cb(tc.key, tc.value)
79110
assert.Equal(t, tc.expKey, k)
80111
assert.Equal(t, tc.expValue, v)
81112
}
113+
if tc.opt != nil {
114+
wantEnv := tc.want.opt.EnvironFunc()
115+
gotEnv := gotProvider.opt.EnvironFunc()
116+
slices.Sort(wantEnv)
117+
slices.Sort(gotEnv)
118+
if !slices.Equal(wantEnv, gotEnv) {
119+
assert.Fail(t, "Env vars not equal (omitted from message for security)",
120+
"Want len: %d\nGot len:%d", len(wantEnv), len(gotEnv))
121+
}
122+
}
82123
})
83124
}
84125
}
85126

86127
func TestProviderWithValue(t *testing.T) {
128+
mockEnviron := func() []string {
129+
return []string{"TEST_FOO=bar"}
130+
}
131+
87132
testCases := []struct {
88-
name string
89-
prefix string
90-
delim string
91-
cb func(key string, value string) (string, interface{})
92-
nilCallback bool
93-
want *Env
133+
name string
134+
prefix string
135+
delim string
136+
cb func(key string, value string) (string, interface{})
137+
opt *Opt
138+
want *Env
94139
}{
95140
{
96-
name: "Nil cb",
97-
prefix: "TEST_",
98-
delim: ".",
99-
nilCallback: true,
141+
name: "Nil cb",
142+
prefix: "TEST_",
143+
delim: ".",
100144
want: &Env{
101145
prefix: "TEST_",
102146
delim: ".",
147+
opt: defaultOpt,
103148
},
104149
},
105150
{
106-
name: "Empty string nil cb",
107-
prefix: "",
108-
delim: ".",
109-
nilCallback: true,
151+
name: "Empty string nil cb",
152+
prefix: "",
153+
delim: ".",
110154
want: &Env{
111155
prefix: "",
112156
delim: ".",
157+
opt: defaultOpt,
113158
},
114159
},
115160
{
@@ -125,6 +170,7 @@ func TestProviderWithValue(t *testing.T) {
125170
cb: func(key string, value string) (string, interface{}) {
126171
return key, value
127172
},
173+
opt: defaultOpt,
128174
},
129175
},
130176
{
@@ -142,23 +188,58 @@ func TestProviderWithValue(t *testing.T) {
142188
key = strings.Replace(strings.TrimPrefix(strings.ToLower(key), "test_"), "_", ".", -1)
143189
return key, value
144190
},
191+
opt: defaultOpt,
192+
},
193+
},
194+
{
195+
name: "Custom opt",
196+
prefix: "TEST_",
197+
delim: ".",
198+
opt: &Opt{
199+
EnvironFunc: mockEnviron,
200+
},
201+
want: &Env{
202+
prefix: "TEST_",
203+
delim: ".",
204+
opt: &Opt{
205+
EnvironFunc: mockEnviron,
206+
},
145207
},
146208
},
147209
}
148210

149211
for _, tc := range testCases {
150212
t.Run(tc.name, func(t *testing.T) {
151-
got := ProviderWithValue(tc.prefix, tc.delim, tc.cb)
152-
if tc.nilCallback {
153-
assert.Equal(t, tc.want, got)
213+
var got *Env
214+
if tc.opt != nil {
215+
got = ProviderWithValue(tc.prefix, tc.delim, tc.cb, tc.opt)
154216
} else {
217+
got = ProviderWithValue(tc.prefix, tc.delim, tc.cb)
218+
}
219+
if tc.cb == nil && tc.opt == nil {
220+
assert.Equal(t, tc.want, got)
221+
return
222+
}
223+
224+
if tc.cb != nil {
155225
keyGot, valGot := got.cb("test_key_env_1", "test_val")
156226
keyWant, valWant := tc.want.cb("test_key_env_1", "test_val")
157227
assert.Equal(t, tc.prefix, got.prefix)
158228
assert.Equal(t, tc.delim, got.delim)
159229
assert.Equal(t, keyWant, keyGot)
160230
assert.Equal(t, valWant, valGot)
161231
}
232+
if tc.opt != nil {
233+
wantEnv := tc.want.opt.EnvironFunc()
234+
gotEnv := got.opt.EnvironFunc()
235+
slices.Sort(wantEnv)
236+
slices.Sort(gotEnv)
237+
if !slices.Equal(wantEnv, gotEnv) {
238+
assert.Fail(t, "Env vars not equal (omitted from message for security)",
239+
"Want len: %d\nGot len:%d", len(wantEnv), len(gotEnv))
240+
}
241+
assert.Equal(t, tc.want.opt.EnvironFunc(), got.opt.EnvironFunc())
242+
}
162243
})
163244
}
164245
}
@@ -180,6 +261,7 @@ func TestRead(t *testing.T) {
180261
expValue: "TEST_VAL",
181262
env: &Env{
182263
delim: ".",
264+
opt: defaultOpt,
183265
},
184266
},
185267
{
@@ -193,6 +275,7 @@ func TestRead(t *testing.T) {
193275
cb: func(key string, value string) (string, interface{}) {
194276
return strings.Replace(strings.ToLower(key), "_", ".", -1), value
195277
},
278+
opt: defaultOpt,
196279
},
197280
},
198281
{
@@ -207,6 +290,7 @@ func TestRead(t *testing.T) {
207290
cb: func(key string, value string) (string, interface{}) {
208291
return strings.Replace(strings.ToLower(key), "_", ".", -1), value
209292
},
293+
opt: defaultOpt,
210294
},
211295
},
212296
{
@@ -217,6 +301,7 @@ func TestRead(t *testing.T) {
217301
expValue: "/test/dir/file",
218302
env: &Env{
219303
delim: ".",
304+
opt: defaultOpt,
220305
},
221306
},
222307
{
@@ -230,6 +315,7 @@ func TestRead(t *testing.T) {
230315
cb: func(key string, value string) (string, interface{}) {
231316
return key, strings.Replace(strings.ToLower(value), "/", "_", -1)
232317
},
318+
opt: defaultOpt,
233319
},
234320
},
235321
{
@@ -240,15 +326,29 @@ func TestRead(t *testing.T) {
240326
expValue: "",
241327
env: &Env{
242328
delim: ".",
329+
opt: defaultOpt,
330+
},
331+
},
332+
{
333+
name: "Environ func provided",
334+
key: "TEST_KEY",
335+
value: "TEST_VAL",
336+
expKey: "TEST_OVERRIDE_KEY",
337+
expValue: "TEST_OVERRIDE_VAL",
338+
env: &Env{
339+
delim: ".",
340+
opt: &Opt{
341+
EnvironFunc: func() []string {
342+
return []string{"TEST_OVERRIDE_KEY=TEST_OVERRIDE_VAL"}
343+
},
344+
},
243345
},
244346
},
245347
}
246348

247349
for _, tc := range testCases {
248350
t.Run(tc.name, func(t *testing.T) {
249-
err := os.Setenv(tc.key, tc.value)
250-
assert.Nil(t, err)
251-
defer os.Unsetenv(tc.key)
351+
t.Setenv(tc.key, tc.value)
252352

253353
envs, err := tc.env.Read()
254354
assert.Nil(t, err)

0 commit comments

Comments
 (0)