Merge pull request aws#1467 from aws/normj/testtool-env-variable-support

Add support for setting environment variables with test tool

Author: normj

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool/LambdaConfigFile.cs

@@ -19,6 +19,9 @@ public class LambdaConfigFile
         [JsonPropertyName("image-command")]
         public string ImageCommand { get; set; }
 
+        [JsonPropertyName("environment-variables")]
+        public string EnvironmentVariables { get; set; }
+
         public string ConfigFileLocation { get; set; }
 
         public string DetermineHandler()

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool/Runtime/LambdaDefaultsConfigFileParser.cs

@@ -6,7 +6,6 @@
 
 using YamlDotNet.RepresentationModel;
 
-
 namespace Amazon.Lambda.TestTool.Runtime
 {
     /// <summary>
@@ -69,6 +68,19 @@ public static LambdaConfigInfo LoadFromFile(LambdaConfigFile configFile)
                 {
                     info.Name = functionHandler;
                 }
+
+                if(configFile.EnvironmentVariables != null)
+                {
+                    var variables = configFile.EnvironmentVariables.Split(';');
+                    foreach(var variable in variables)
+                    {
+                        var kvp = variable.Split('=');
+                        if(kvp.Length == 2)
+                        {
+                            info.EnvironmentVariables[kvp[0]] = kvp[1];
+                        }
+                    }
+                }
 
                 configInfo.FunctionInfos.Add(info);
             }
@@ -197,6 +209,21 @@ private static void ProcessYamlServerlessTemplateFunctionBased(LambdaConfigInfo
                     Handler = handler
                 };
 
+                if (resourceBody.Children.TryGetValue("Environment", out var environmentProperty) && environmentProperty is YamlMappingNode)
+                {
+                    if (((YamlMappingNode)environmentProperty).Children.TryGetValue("Environment", out var variableProperty) && variableProperty is YamlMappingNode)
+                    {
+                        foreach(var kvp in ((YamlMappingNode)variableProperty).Children) 
+                        {
+                            if(kvp.Key is YamlScalarNode keyNode && keyNode.Value != null && 
+                                kvp.Value is YamlScalarNode valueNode && valueNode.Value != null)
+                            {
+                                functionInfo.EnvironmentVariables[keyNode.Value] = valueNode.Value;
+                            }
+                        }
+                    }
+                }
+
                 configInfo.FunctionInfos.Add(functionInfo);
             }
         }
@@ -255,6 +282,18 @@ private static void ProcessJsonServerlessTemplate(LambdaConfigInfo configInfo, s
                         Handler = handler
                     };
 
+                    if(propertiesProperty.TryGetProperty("Environment", out var environmentProperty) && 
+                        environmentProperty.TryGetProperty("Variables", out var variablesProperty))
+                    {
+                        foreach(var property in variablesProperty.EnumerateObject())
+                        {
+                            if(property.Value.ValueKind == JsonValueKind.String)
+                            {
+                                functionInfo.EnvironmentVariables[property.Name] = property.Value.GetString();
+                            }
+                        }
+                    }
+
                     configInfo.FunctionInfos.Add(functionInfo);
                 }
             }

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool/Runtime/LambdaExecutor.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
 using System.Linq;
@@ -37,6 +38,18 @@ public async Task<ExecutionResponse> ExecuteAsync(ExecutionRequest request)
                     Environment.SetEnvironmentVariable("AWS_PROFILE", request.AWSProfile);
                 }
 
+                // Set the Lambda environment variable for the function name. Some libraries like
+                // our Amazon.Lambda.AspNetCoreServer.Hosting use this environment variable to
+                // tell if they are running in Lambda and if so activate. Since we are emulating
+                // Lambda we want those libraries to activate.
+                Environment.SetEnvironmentVariable("AWS_LAMBDA_FUNCTION_NAME", request.Function.FunctionInfo.Name);
+
+                // If Environment variables were defined for the function
+                // then set them for the process to the emulated function picks up the variables.
+                foreach (var kvp in request.Function.FunctionInfo.EnvironmentVariables)
+                {
+                    Environment.SetEnvironmentVariable(kvp.Key, kvp.Value);
+                }
 
                 var context = new LocalLambdaContext()
                 {
@@ -64,6 +77,12 @@ public async Task<ExecutionResponse> ExecuteAsync(ExecutionRequest request)
                 }
                 finally
                 {
+                    // To avoid side effects remove the environment variables that were set specifically 
+                    // for running the lambda function.
+                    foreach (var kvp in request.Function.FunctionInfo.EnvironmentVariables)
+                    {
+                        Environment.SetEnvironmentVariable(kvp.Key, null);
+                    }
                     executeSlim.Release();
                 }
             }

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool/Runtime/LambdaFunctionInfo.cs

@@ -1,4 +1,6 @@
-namespace Amazon.Lambda.TestTool.Runtime
+using System.Collections.Generic;
+
+namespace Amazon.Lambda.TestTool.Runtime
 {
     public class LambdaFunctionInfo
     {
@@ -11,5 +13,7 @@ public class LambdaFunctionInfo
         /// The Lambda function handler string.
         /// </summary>
         public string Handler { get; set; }
+
+        public IDictionary<string, string> EnvironmentVariables { get; }  = new Dictionary<string, string>();
     }
 }

Tools/LambdaTestTool/tests/Amazon.Lambda.TestTool.Tests.Shared/DefaultsFileParseTests.cs

@@ -41,6 +41,28 @@ public void LambdaFunctionWithName()
                 File.Delete(jsonFile);
             }
         }
+
+        [Fact]
+        public void LambdaFunctionWithEnvironmentVariables()
+        {
+            var jsonFile = WriteTempConfigFile("{\"function-handler\" : \"Assembly::Type::Method\", \"environment-variables\" : \"key1=value1;key2=value2\"}");
+            try
+            {
+                var configInfo = LambdaDefaultsConfigFileParser.LoadFromFile(jsonFile);
+                Assert.Single(configInfo.FunctionInfos);
+                Assert.Equal("Assembly::Type::Method", configInfo.FunctionInfos[0].Handler);
+                Assert.Equal("TheFunc", configInfo.FunctionInfos[0].Name);
+
+                Assert.Equal(2, configInfo.FunctionInfos[0].EnvironmentVariables.Count);
+                Assert.Equal("value1", configInfo.FunctionInfos[0].EnvironmentVariables["key1"]);
+                Assert.Equal("value2", configInfo.FunctionInfos[0].EnvironmentVariables["key2"]);
+            }
+            finally
+            {
+                File.Delete(jsonFile);
+            }
+        }
+
         [Fact]
         public void LambdaFunctionWithImageCommand()
         {