diff --git a/execute_lambda.py b/execute_lambda.py
index b4cbb4a53de..b599b6636ca 100644
--- a/execute_lambda.py
+++ b/execute_lambda.py
@@ -260,8 +260,10 @@ def main():
 
 def wait_for_lambda(client, module, name):
     try:
-        waiter = client.get_waiter('function_active')
-        waiter.wait(FunctionName=name)
+        client_active_waiter = client.get_waiter('function_active')
+        client_updated_waiter = client.get_waiter('function_updated')
+        client_active_waiter.wait(FunctionName=name)
+        client_updated_waiter.wait(FunctionName=name)
     except botocore.exceptions.WaiterError as e:
         module.fail_json_aws(e, msg='Timeout while waiting on lambda to be Active')
     except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
diff --git a/lambda.py b/lambda.py
index 8ad45785f21..ff469c5bc71 100644
--- a/lambda.py
+++ b/lambda.py
@@ -107,6 +107,11 @@
     description:
       - Tag dict to apply to the function.
     type: dict
+  kms_key_arn:
+    description:
+      - The KMS key ARN used to encrypt the function's environment variables.
+    type: str
+    version_added: 3.3.0
 author:
     - 'Steyn Huizinga (@steynovich)'
 extends_documentation_fragment:
@@ -451,6 +456,7 @@ def main():
         vpc_security_group_ids=dict(type='list', elements='str'),
         environment_variables=dict(type='dict'),
         dead_letter_arn=dict(),
+        kms_key_arn=dict(type='str', no_log=False),
         tracing_mode=dict(choices=['Active', 'PassThrough']),
         tags=dict(type='dict'),
     )
@@ -488,6 +494,7 @@ def main():
     dead_letter_arn = module.params.get('dead_letter_arn')
     tracing_mode = module.params.get('tracing_mode')
     tags = module.params.get('tags')
+    kms_key_arn = module.params.get('kms_key_arn')
 
     check_mode = module.check_mode
     changed = False
@@ -543,6 +550,8 @@ def main():
                     func_kwargs.update({'DeadLetterConfig': {'TargetArn': dead_letter_arn}})
         if tracing_mode and (current_config.get('TracingConfig', {}).get('Mode', 'PassThrough') != tracing_mode):
             func_kwargs.update({'TracingConfig': {'Mode': tracing_mode}})
+        if kms_key_arn:
+            func_kwargs.update({'KMSKeyArn': kms_key_arn})
 
         # If VPC configuration is desired
         if vpc_subnet_ids:
@@ -674,17 +683,23 @@ def main():
         if tracing_mode:
             func_kwargs.update({'TracingConfig': {'Mode': tracing_mode}})
 
+        if kms_key_arn:
+            func_kwargs.update({'KMSKeyArn': kms_key_arn})
+
         # If VPC configuration is given
         if vpc_subnet_ids:
             func_kwargs.update({'VpcConfig': {'SubnetIds': vpc_subnet_ids,
                                               'SecurityGroupIds': vpc_security_group_ids}})
 
+        # Function would have been created if not check mode
+        if check_mode:
+            module.exit_json(changed=True)
+
         # Finally try to create function
         current_version = None
         try:
-            if not check_mode:
-                response = client.create_function(aws_retry=True, **func_kwargs)
-                current_version = response['Version']
+            response = client.create_function(aws_retry=True, **func_kwargs)
+            current_version = response['Version']
             changed = True
         except (BotoCoreError, ClientError) as e:
             module.fail_json_aws(e, msg="Trying to create function")