Add Merge method to Environment

Signed-off-by: Burak Varlı <[email protected]>
awslabs · Jan 30, 2025 · ede6473 · ede6473
1 parent 8adcb20
commit ede6473
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 0 deletions.
diff --git a/pkg/driver/node/envprovider/provider.go b/pkg/driver/node/envprovider/provider.go
@@ -83,6 +83,14 @@ func (env Environment) Set(key Key, value Value) {
 	env[key] = value
 }
 
+// Merge adds all key-value pairs from the given environment to the current environment.
+// If a key exists in both environments, the value from the given environment takes precedence.
+func (env Environment) Merge(other Environment) {
+	for key, value := range other {
+		env[key] = value
+	}
+}
+
 // format formats given key and value to be used as an environment variable.
 func format(key Key, value Value) string {
 	return fmt.Sprintf("%s=%s", key, value)

diff --git a/pkg/driver/node/envprovider/provider_test.go b/pkg/driver/node/envprovider/provider_test.go
@@ -219,3 +219,50 @@ func TestEnvironmentList(t *testing.T) {
 		})
 	}
 }
+
+func TestMergingEnvironments(t *testing.T) {
+	testCases := []struct {
+		name  string
+		env   envprovider.Environment
+		other envprovider.Environment
+		want  envprovider.Environment
+	}{
+		{
+			name:  "merge into empty environment",
+			env:   envprovider.Environment{},
+			other: envprovider.Environment{"AWS_REGION": "us-west-1"},
+			want:  envprovider.Environment{"AWS_REGION": "us-west-1"},
+		},
+		{
+			name:  "merge empty environment",
+			env:   envprovider.Environment{"AWS_REGION": "us-west-1"},
+			other: envprovider.Environment{},
+			want:  envprovider.Environment{"AWS_REGION": "us-west-1"},
+		},
+		{
+			name:  "merge with different keys",
+			env:   envprovider.Environment{"AWS_REGION": "us-west-1"},
+			other: envprovider.Environment{"AWS_DEFAULT_REGION": "us-east-1"},
+			want:  envprovider.Environment{"AWS_REGION": "us-west-1", "AWS_DEFAULT_REGION": "us-east-1"},
+		},
+		{
+			name:  "merge with overlapping keys",
+			env:   envprovider.Environment{"AWS_REGION": "us-west-1", "AWS_PROFILE": "default"},
+			other: envprovider.Environment{"AWS_REGION": "us-east-1", "AWS_DEFAULT_REGION": "us-east-2"},
+			want:  envprovider.Environment{"AWS_REGION": "us-east-1", "AWS_PROFILE": "default", "AWS_DEFAULT_REGION": "us-east-2"},
+		},
+		{
+			name:  "merge with empty values",
+			env:   envprovider.Environment{"AWS_REGION": "us-west-1"},
+			other: envprovider.Environment{"AWS_REGION": "", "AWS_DEFAULT_REGION": "us-east-1"},
+			want:  envprovider.Environment{"AWS_REGION": "", "AWS_DEFAULT_REGION": "us-east-1"},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			testCase.env.Merge(testCase.other)
+			assert.Equals(t, testCase.want, testCase.env)
+		})
+	}
+}