From 9e55d2ee9803bc0d3935f65247ea29ca9cbcbf9c Mon Sep 17 00:00:00 2001
From: Onur Filiz <onfiliz@amazon.com>
Date: Fri, 3 Aug 2018 11:49:29 -0700
Subject: [PATCH] Fargate: Add environment variable support

---
 providers/aws/fargate/container.go      | 12 ++++++++++++
 providers/aws/fargate/container_test.go | 15 ++++++++++++++-
 providers/aws/provider_test.go          | 22 +++++++++++++++++++---
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/providers/aws/fargate/container.go b/providers/aws/fargate/container.go
index de9ccb9b3..5e16a853c 100644
--- a/providers/aws/fargate/container.go
+++ b/providers/aws/fargate/container.go
@@ -52,6 +52,18 @@ func newContainer(spec *corev1.Container) (*container, error) {
 		cntr.definition.WorkingDirectory = aws.String(spec.WorkingDir)
 	}
 
+	// Add environment variables.
+	if spec.Env != nil {
+		for _, env := range spec.Env {
+			cntr.definition.Environment = append(
+				cntr.definition.Environment,
+				&ecs.KeyValuePair{
+					Name:  aws.String(env.Name),
+					Value: aws.String(env.Value),
+				})
+		}
+	}
+
 	// Translate the Kubernetes container resource requirements to Fargate units.
 	cntr.setResourceRequirements(&spec.Resources)
 
diff --git a/providers/aws/fargate/container_test.go b/providers/aws/fargate/container_test.go
index a563049db..5700885d3 100644
--- a/providers/aws/fargate/container_test.go
+++ b/providers/aws/fargate/container_test.go
@@ -32,6 +32,10 @@ var (
 		Command:    []string{"anyCmd"},
 		Args:       []string{"anyArg1", "anyArg2"},
 		WorkingDir: "/any/working/dir",
+		Env: []corev1.EnvVar{
+			{Name: "anyEnvName1", Value: "anyEnvValue1"},
+			{Name: "anyEnvName2", Value: "anyEnvValue2"},
+		},
 	}
 )
 
@@ -46,8 +50,17 @@ func TestContainerDefinition(t *testing.T) {
 	assert.Equal(t, cntrSpec.Name, *cntr.definition.Name, "incorrect name")
 	assert.Equal(t, cntrSpec.Image, *cntr.definition.Image, "incorrect image")
 	assert.Equal(t, cntrSpec.Command[0], *cntr.definition.EntryPoint[0], "incorrect command")
-	assert.Equal(t, cntrSpec.Args[0], *cntr.definition.Command[0], "incorrect args")
+
+	for i, env := range cntrSpec.Args {
+		assert.Equal(t, env, *cntr.definition.Command[i], "incorrect args")
+	}
+
 	assert.Equal(t, cntrSpec.WorkingDir, *cntr.definition.WorkingDirectory, "incorrect working dir")
+
+	for i, env := range cntrSpec.Env {
+		assert.Equal(t, env.Name, *cntr.definition.Environment[i].Name, "incorrect env name")
+		assert.Equal(t, env.Value, *cntr.definition.Environment[i].Value, "incorrect env value")
+	}
 }
 
 // TestContainerResourceRequirementsDefaults verifies whether the container gets default CPU
diff --git a/providers/aws/provider_test.go b/providers/aws/provider_test.go
index 79d10d90a..b38a9f0e6 100644
--- a/providers/aws/provider_test.go
+++ b/providers/aws/provider_test.go
@@ -237,7 +237,13 @@ func TestAWSFargateProviderPodLifecycle(t *testing.T) {
 						"/bin/sh",
 					},
 					Args: []string{
-						"-c", "echo \"Started\"; while true; do sleep 1; done",
+						"-c",
+						"echo \"Started\";" +
+							"echo \"TEST_ENV=$TEST_ENV\";" +
+							"while true; do sleep 1; done",
+					},
+					Env: []v1.EnvVar{
+						{Name: "TEST_ENV", Value: "AnyValue"},
 					},
 					Resources: v1.ResourceRequirements{
 						Limits: v1.ResourceList{
@@ -281,8 +287,18 @@ func TestAWSFargateProviderPodLifecycle(t *testing.T) {
 			t.Error(err)
 		}
 
-		if logs != "Started\n" {
-			t.Errorf("Expected logs to be \"Started\\n\", but received \"%v\"", logs)
+		// Test log output.
+		receivedLogs := strings.Split(logs, "\n")
+		expectedLogs := []string{
+			"Started",
+			pod.Spec.Containers[0].Env[0].Name + "=" + pod.Spec.Containers[0].Env[0].Value,
+		}
+
+		for i, line := range receivedLogs {
+			fmt.Printf("Log[#%d]: %v\n", i, line)
+			if len(expectedLogs) > i && receivedLogs[i] != expectedLogs[i] {
+				t.Errorf("Expected log line %d to be %q, but received %q", i, line, receivedLogs[i])
+			}
 		}
 
 		// Delete the pod.