Repository: getmoto/moto
Pull request #7497 (2024-03-20 09:54:23)

Lambda: simple lambda custom responses
Customise Lambda response by posting desired results to a new, static moto endpoint.

Resolves https://github.com/getmoto/moto/issues/7481

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/getmoto/moto/pull/7497?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 94.05%. Comparing base [(`9aef694`)](https://app.codecov.io/gh/getmoto/moto/commit/9aef6944da6a6f1b84be99cefb34f1724e6ca58f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) to head [(`7fd6ff7`)](https://app.codecov.io/gh/getmoto/moto/pull/7497?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).
> Report is 122 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #7497      +/-   ##
==========================================
- Coverage   95.88%   94.05%   -1.83%     
==========================================
  Files         843     1064     +221     
  Lines       82578    91444    +8866     
==========================================
+ Hits        79178    86007    +6829     
- Misses       3400     5437    +2037     
```

| [Flag](https://app.codecov.io/gh/getmoto/moto/pull/7497/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | Coverage Δ | |
|---|---|---|
| [servertests](https://app.codecov.io/gh/getmoto/moto/pull/7497/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `29.65% <19.04%> (-2.88%)` | :arrow_down: |
| [unittests](https://app.codecov.io/gh/getmoto/moto/pull/7497/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `94.02% <100.00%> (-1.83%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto#carryforward-flags-in-the-pull-request-comment) to find out more.


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/getmoto/moto/pull/7497?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).


Patch:
diff --git a/docs/docs/services/lambda.rst b/docs/docs/services/lambda.rst
index dbbad8de8..a4e91e2ac 100644
--- a/docs/docs/services/lambda.rst
+++ b/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
diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py
index 3951c0c76..c7a12d6a4 100644
--- a/moto/awslambda/responses.py
+++ b/moto/awslambda/responses.py
@@ -255,7 +255,7 @@ class LambdaResponse(BaseResponse):
         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"
diff --git a/moto/awslambda_simple/models.py b/moto/awslambda_simple/models.py
index 2fa786632..1be101c8f 100644
--- a/moto/awslambda_simple/models.py
+++ b/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 @@ class LambdaSimpleBackend(LambdaBackend):
         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")
diff --git a/moto/moto_api/_internal/models.py b/moto/moto_api/_internal/models.py
index d490f6010..1b00d3d5f 100644
--- a/moto/moto_api/_internal/models.py
+++ b/moto/moto_api/_internal/models.py
@@ -54,6 +54,14 @@ class MotoAPIBackend(BaseBackend):
         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,
diff --git a/moto/moto_api/_internal/responses.py b/moto/moto_api/_internal/responses.py
index 665b28997..1aceb29e4 100644
--- a/moto/moto_api/_internal/responses.py
+++ b/moto/moto_api/_internal/responses.py
@@ -179,6 +179,24 @@ class MotoAPIResponse(BaseResponse):
             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,
diff --git a/moto/moto_api/_internal/urls.py b/moto/moto_api/_internal/urls.py
index e0b53b794..278ed7249 100644
--- a/moto/moto_api/_internal/urls.py
+++ b/moto/moto_api/_internal/urls.py
@@ -17,6 +17,7 @@ url_paths = {
     "{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,


Test patch:
diff --git a/tests/test_awslambda_simple/test_lambda_simple.py b/tests/test_awslambda_simple/test_lambda_simple.py
index 3d6d78e86..be4454c28 100644
--- a/tests/test_awslambda_simple/test_lambda_simple.py
+++ b/tests/test_awslambda_simple/test_lambda_simple.py
@@ -2,6 +2,7 @@ import json
 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():
diff --git a/tests/test_stepfunctions/parser/test_stepfunctions_sns_integration.py b/tests/test_stepfunctions/parser/test_stepfunctions_sns_integration.py
index 26fb4b79b..54e3a61c1 100644
--- a/tests/test_stepfunctions/parser/test_stepfunctions_sns_integration.py
+++ b/tests/test_stepfunctions/parser/test_stepfunctions_sns_integration.py
@@ -33,4 +33,6 @@ def test_state_machine_calling_sns_publish():
         _, 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)
+    )