Lambda: simple lambda custom responses (#7497)

docs/docs/services/lambda.rst

@@ -69,7 +69,27 @@ lambda
 - [X] invoke
 
         Invoking a Function with PackageType=Image is not yet supported.
-
+
+        Invoking a Funcation against Lambda without docker now supports customised responses, the default being `Simple Lambda happy path OK`.
+        You can use a dedicated API to override this, by configuring a queue of expected results.
+
+        A request to `invoke` will take the first result from that queue.
+
+        Configure this queue by making an HTTP request to `/moto-api/static/lambda-simple/response`. An example invocation looks like this:
+
+        .. sourcecode:: python
+
+            expected_results = {"results": ["test", "test 2"], "region": "us-east-1"}
+            resp = requests.post(
+                    "http://motoapi.amazonaws.com:5000/moto-api/static/lambda-simple/response",
+                    json=expected_results,
+                )
+
+            assert resp.status_code == 201
+
+            client = boto3.client("lambda", region_name="us-east-1")
+            resp = client.invoke(...) # resp["Payload"].read().decode() == "test"
+            resp = client.invoke(...) # resp["Payload"].read().decode() == "test2"
 
 - [ ] invoke_async
 - [ ] invoke_with_response_stream

moto/awslambda/responses.py

@@ -255,7 +255,7 @@ def _invoke(self, request: Any) -> Tuple[int, Dict[str, str], Union[str, bytes]]
         payload = self.backend.invoke(
             function_name, qualifier, self.body, self.headers, response_headers
         )
-        if payload:
+        if payload is not None:
             if request.headers.get("X-Amz-Invocation-Type") != "Event":
                 if sys.getsizeof(payload) > 6000000:
                     response_headers["Content-Length"] = "142"

moto/awslambda_simple/models.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union
+from typing import Any, List, Optional, Union
 
 from moto.awslambda.models import LambdaBackend
 from moto.core.base_backend import BackendDict
@@ -10,6 +10,10 @@ class LambdaSimpleBackend(LambdaBackend):
     Annotate your tests with `@mock_aws(config={"lambda": {"use_docker": False}}) to use this Lambda-implementation.
     """
 
+    def __init__(self, region_name: str, account_id: str):
+        super().__init__(region_name, account_id)
+        self.lambda_simple_results_queue: List[str] = []
+
     # pylint: disable=unused-argument
     def invoke(
         self,
@@ -20,9 +24,10 @@ def invoke(
         response_headers: Any,
     ) -> Optional[Union[str, bytes]]:
 
-        if body:
-            return str.encode(body)
-        return b"Simple Lambda happy path OK"
+        default_result = "Simple Lambda happy path OK"
+        if self.lambda_simple_results_queue:
+            default_result = self.lambda_simple_results_queue.pop(0)
+        return str.encode(default_result)
 
 
 lambda_simple_backends = BackendDict(LambdaSimpleBackend, "lambda")

moto/moto_api/_internal/models.py

@@ -54,6 +54,14 @@ def set_ce_cost_usage(self, result: Dict[str, Any], account_id: str) -> None:
         backend = ce_backends[account_id]["global"]
         backend.cost_usage_results_queue.append(result)
 
+    def set_lambda_simple_result(
+        self, result: str, account_id: str, region: str
+    ) -> None:
+        from moto.awslambda_simple.models import lambda_simple_backends
+
+        backend = lambda_simple_backends[account_id][region]
+        backend.lambda_simple_results_queue.append(result)
+
     def set_sagemaker_result(
         self,
         body: str,

moto/moto_api/_internal/responses.py

@@ -179,6 +179,24 @@ def set_ce_cost_usage_result(
             moto_api_backend.set_ce_cost_usage(result=result, account_id=account_id)
         return 201, {}, ""
 
+    def set_lambda_simple_result(
+        self,
+        request: Any,
+        full_url: str,  # pylint: disable=unused-argument
+        headers: Any,
+    ) -> TYPE_RESPONSE:
+        from .models import moto_api_backend
+
+        body = self._get_body(headers, request)
+        account_id = body.get("account_id", DEFAULT_ACCOUNT_ID)
+        region = body.get("region", "us-east-1")
+
+        for result in body.get("results", []):
+            moto_api_backend.set_lambda_simple_result(
+                result=result, account_id=account_id, region=region
+            )
+        return 201, {}, ""
+
     def set_sagemaker_result(
         self,
         request: Any,

moto/moto_api/_internal/urls.py

@@ -17,6 +17,7 @@
     "{0}/moto-api/static/athena/query-results": response_instance.set_athena_result,
     "{0}/moto-api/static/ce/cost-and-usage-results": response_instance.set_ce_cost_usage_result,
     "{0}/moto-api/static/inspector2/findings-results": response_instance.set_inspector2_findings_result,
+    "{0}/moto-api/static/lambda-simple/response": response_instance.set_lambda_simple_result,
     "{0}/moto-api/static/sagemaker/endpoint-results": response_instance.set_sagemaker_result,
     "{0}/moto-api/static/rds-data/statement-results": response_instance.set_rds_data_result,
     "{0}/moto-api/state-manager/get-transition": response_instance.get_transition,

tests/test_awslambda_simple/test_lambda_simple.py

@@ -2,6 +2,7 @@
 from unittest import SkipTest
 
 import boto3
+import requests
 
 from moto import mock_aws, settings
 
@@ -42,7 +43,44 @@ def test_run_function_no_log():
 
     # Verify
     assert result["StatusCode"] == 200
-    assert json.loads(result["Payload"].read().decode("utf-8")) == payload
+    assert result["Payload"].read().decode("utf-8") == "Simple Lambda happy path OK"
+
+
+@mock_aws(config={"lambda": {"use_docker": False}})
+def test_set_lambda_simple_query_results():
+    # Setup
+    base_url = (
+        settings.test_server_mode_endpoint()
+        if settings.TEST_SERVER_MODE
+        else "http://motoapi.amazonaws.com"
+    )
+    results = {"results": ["test", "test 2"], "region": LAMBDA_REGION}
+    resp = requests.post(
+        f"{base_url}/moto-api/static/lambda-simple/response",
+        json=results,
+    )
+    assert resp.status_code == 201
+
+    client = setup_lambda()
+
+    # Execute & Verify
+    resp = client.invoke(
+        FunctionName=FUNCTION_NAME,
+        LogType="Tail",
+    )
+    assert resp["Payload"].read().decode() == results["results"][0]
+
+    resp = client.invoke(
+        FunctionName=FUNCTION_NAME,
+        LogType="Tail",
+    )
+    assert resp["Payload"].read().decode() == results["results"][1]
+
+    resp = client.invoke(
+        FunctionName=FUNCTION_NAME,
+        LogType="Tail",
+    )
+    assert resp["Payload"].read().decode() == "Simple Lambda happy path OK"
 
 
 def setup_lambda():

tests/test_stepfunctions/parser/test_stepfunctions_sns_integration.py

@@ -33,4 +33,6 @@ def _verify_result(client, execution, execution_arn):
         _, msg, _, _, _ = notification
         assert msg == "my msg"
 
-    verify_execution_result(_verify_result, expected_status, tmpl_name, exec_input=json.dumps(exec_input))
+    verify_execution_result(
+        _verify_result, expected_status, tmpl_name, exec_input=json.dumps(exec_input)
+    )