From 0cd969c9c14179b7e79d34b61d7a43f2bfbdd93e Mon Sep 17 00:00:00 2001
From: Vincent <97131062+vincbeck@users.noreply.github.com>
Date: Fri, 2 Dec 2022 14:18:47 -0500
Subject: [PATCH] Lambda hook: make runtime and handler optional (#27778)

* Lambda hook: make runtime and handler optional
---
 .../amazon/aws/hooks/lambda_function.py       |  10 +-
 .../amazon/aws/hooks/test_lambda_function.py  | 108 +++++++++++++++---
 2 files changed, 103 insertions(+), 15 deletions(-)

diff --git a/airflow/providers/amazon/aws/hooks/lambda_function.py b/airflow/providers/amazon/aws/hooks/lambda_function.py
index 2919d37701aed..2df2fb66207ab 100644
--- a/airflow/providers/amazon/aws/hooks/lambda_function.py
+++ b/airflow/providers/amazon/aws/hooks/lambda_function.py
@@ -72,9 +72,9 @@ def create_lambda(
         self,
         *,
         function_name: str,
-        runtime: str,
+        runtime: str | None = None,
         role: str,
-        handler: str,
+        handler: str | None = None,
         code: dict,
         description: str | None = None,
         timeout: int | None = None,
@@ -93,6 +93,12 @@ def create_lambda(
         code_signing_config_arn: str | None = None,
         architectures: list[str] | None = None,
     ) -> dict:
+        if package_type == "Zip":
+            if handler is None:
+                raise TypeError("Parameter 'handler' is required if 'package_type' is 'Zip'")
+            if runtime is None:
+                raise TypeError("Parameter 'runtime' is required if 'package_type' is 'Zip'")
+
         """Create a Lambda Function"""
         create_function_args = {
             "FunctionName": function_name,
diff --git a/tests/providers/amazon/aws/hooks/test_lambda_function.py b/tests/providers/amazon/aws/hooks/test_lambda_function.py
index b6f1edaa12310..f2d9ca5ee30d3 100644
--- a/tests/providers/amazon/aws/hooks/test_lambda_function.py
+++ b/tests/providers/amazon/aws/hooks/test_lambda_function.py
@@ -17,31 +17,113 @@
 # under the License.
 from __future__ import annotations
 
-from unittest.mock import patch
+from unittest import mock
 
-from moto import mock_lambda
+import pytest
 
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
 
+FUNCTION_NAME = "test_function"
+PAYLOAD = '{"hello": "airflow"}'
+RUNTIME = "python3.9"
+ROLE = "role"
+HANDLER = "handler"
+CODE = {}
+
 
-@mock_lambda
 class TestLambdaHook:
     def test_get_conn_returns_a_boto3_connection(self):
         hook = LambdaHook(aws_conn_id="aws_default")
         assert hook.conn is not None
 
-    def test_invoke_lambda_function(self):
+    @mock.patch(
+        "airflow.providers.amazon.aws.hooks.lambda_function.LambdaHook.conn", new_callable=mock.PropertyMock
+    )
+    def test_invoke_lambda(self, mock_conn):
+        mock_conn().invoke.return_value = {}
+
+        hook = LambdaHook(aws_conn_id="aws_default")
+        hook.invoke_lambda(function_name=FUNCTION_NAME, payload=PAYLOAD)
+
+        mock_conn().invoke.assert_called_once_with(
+            FunctionName=FUNCTION_NAME,
+            Payload=PAYLOAD,
+        )
+
+    @mock.patch(
+        "airflow.providers.amazon.aws.hooks.lambda_function.LambdaHook.conn", new_callable=mock.PropertyMock
+    )
+    def test_create_lambda_with_zip_package_type(self, mock_conn):
+        mock_conn().create_function.return_value = {}
+
+        hook = LambdaHook(aws_conn_id="aws_default")
+        hook.create_lambda(
+            function_name=FUNCTION_NAME,
+            runtime=RUNTIME,
+            role=ROLE,
+            handler=HANDLER,
+            code=CODE,
+            package_type="Zip",
+        )
+
+        mock_conn().create_function.assert_called_once_with(
+            FunctionName=FUNCTION_NAME,
+            Runtime=RUNTIME,
+            Role=ROLE,
+            Handler=HANDLER,
+            Code=CODE,
+            PackageType="Zip",
+        )
+
+    @mock.patch(
+        "airflow.providers.amazon.aws.hooks.lambda_function.LambdaHook.conn", new_callable=mock.PropertyMock
+    )
+    def test_create_lambda_with_zip_package_type_and_no_runtime(self, mock_conn):
+        mock_conn().create_function.return_value = {}
 
         hook = LambdaHook(aws_conn_id="aws_default")
+        with pytest.raises(TypeError):
+            hook.create_lambda(
+                function_name=FUNCTION_NAME,
+                role=ROLE,
+                handler=HANDLER,
+                code=CODE,
+                package_type="Zip",
+            )
 
-        with patch.object(hook.conn, "invoke") as mock_invoke:
-            payload = '{"hello": "airflow"}'
-            hook.invoke_lambda(function_name="test_function", payload=payload)
+    @mock.patch(
+        "airflow.providers.amazon.aws.hooks.lambda_function.LambdaHook.conn", new_callable=mock.PropertyMock
+    )
+    def test_create_lambda_with_zip_package_type_and_no_handler(self, mock_conn):
+        mock_conn().create_function.return_value = {}
+
+        hook = LambdaHook(aws_conn_id="aws_default")
+        with pytest.raises(TypeError):
+            hook.create_lambda(
+                function_name=FUNCTION_NAME,
+                runtime=RUNTIME,
+                role=ROLE,
+                code=CODE,
+                package_type="Zip",
+            )
+
+    @mock.patch(
+        "airflow.providers.amazon.aws.hooks.lambda_function.LambdaHook.conn", new_callable=mock.PropertyMock
+    )
+    def test_create_lambda_with_image_package_type(self, mock_conn):
+        mock_conn().create_function.return_value = {}
+
+        hook = LambdaHook(aws_conn_id="aws_default")
+        hook.create_lambda(
+            function_name=FUNCTION_NAME,
+            role=ROLE,
+            code=CODE,
+            package_type="Image",
+        )
 
-        mock_invoke.asset_called_once_with(
-            FunctionName="test_function",
-            InvocationType="RequestResponse",
-            LogType="None",
-            Payload=payload,
-            Qualifier="$LATEST",
+        mock_conn().create_function.assert_called_once_with(
+            FunctionName=FUNCTION_NAME,
+            Role=ROLE,
+            Code=CODE,
+            PackageType="Image",
         )