From b98fca5b14675b860c9f13c140a420605ffaa250 Mon Sep 17 00:00:00 2001
From: Simon Tabor <me@simontabor.com>
Date: Tue, 11 Aug 2020 21:04:12 +0100
Subject: [PATCH] awslambda put_function_env functionality

Signed-off-by: Simon Tabor <me@simontabor.com>
---
 chaosaws/awslambda/actions.py             | 17 ++++++++++++++++-
 tests/awslambda/test_awslambda_actions.py | 15 ++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/chaosaws/awslambda/actions.py b/chaosaws/awslambda/actions.py
index 6cf2b90..cb4a8a8 100644
--- a/chaosaws/awslambda/actions.py
+++ b/chaosaws/awslambda/actions.py
@@ -124,13 +124,26 @@ def put_function_memory_size(function_name: str, memory_size: int,
                                           memory_size=memory_size)
 
 
+def put_function_env(function_name: str, env: Dict[str, str],
+                     configuration: Configuration = None,
+                     secrets: Secrets = None) -> AWSResponse:
+    """
+    Sets environment variables for the function.
+    """
+    client = aws_client("lambda", configuration, secrets)
+    return _update_function_configuration(function_name=function_name,
+                                          client=client,
+                                          env=env)
+
+
 ###############################################################################
 # Private functions
 ###############################################################################
 def _update_function_configuration(function_name: str,
                                    client: boto3.client,
                                    timeout: int = None,
-                                   memory_size: int = None) -> AWSResponse:
+                                   memory_size: int = None,
+                                   env: Dict[str, str] = None) -> AWSResponse:
     request_kwargs = {
         'FunctionName': function_name
     }
@@ -138,4 +151,6 @@ def _update_function_configuration(function_name: str,
         request_kwargs['Timeout'] = timeout
     if memory_size is not None:
         request_kwargs['MemorySize'] = memory_size
+    if env is not None:
+        request_kwargs['Environment'] = {'Variables': env}
     return client.update_function_configuration(**request_kwargs)
diff --git a/tests/awslambda/test_awslambda_actions.py b/tests/awslambda/test_awslambda_actions.py
index 78ca859..b17da07 100644
--- a/tests/awslambda/test_awslambda_actions.py
+++ b/tests/awslambda/test_awslambda_actions.py
@@ -10,7 +10,8 @@
                                         invoke_function,
                                         put_function_concurrency,
                                         put_function_memory_size,
-                                        put_function_timeout)
+                                        put_function_timeout,
+                                        put_function_env)
 
 
 @patch('chaosaws.awslambda.actions.aws_client', autospec=True)
@@ -166,3 +167,15 @@ def test_aws_lambda_put_function_memory_size(aws_client):
     client.update_function_configuration.assert_called_with(
         FunctionName=lambda_function_name,
         MemorySize=memory_size)
+
+
+@patch('chaosaws.awslambda.actions.aws_client', autospec=True)
+def test_aws_lambda_put_function_env(aws_client):
+    client = MagicMock()
+    aws_client.return_value = client
+    lambda_function_name = 'my-lambda-function'
+    env = {'Testing': 'true'}
+    put_function_env(lambda_function_name, env)
+    client.update_function_configuration.assert_called_with(
+        FunctionName=lambda_function_name,
+        Environment={'Variables': env})