From a616f1701a3617d8cc75063c9caae6052c3e1699 Mon Sep 17 00:00:00 2001
From: Khoi Pham <dkpham@amazon.com>
Date: Fri, 2 Aug 2024 10:33:18 -0700
Subject: [PATCH] fix: Remove str type constrain in Fn:FindInMap to avoid the
 issue of lookup int value

---
 .../intrinsic_property_resolver.py            |  2 +-
 .../local/start_lambda/test_start_lambda.py   | 22 +++++++++++++++++++
 .../integration/testdata/invoke/template.yml  | 13 +++++++++++
 .../test_intrinsic_resolver.py                |  6 +++++
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/samcli/lib/intrinsic_resolver/intrinsic_property_resolver.py b/samcli/lib/intrinsic_resolver/intrinsic_property_resolver.py
index 1fee887de7..667b222a93 100644
--- a/samcli/lib/intrinsic_resolver/intrinsic_property_resolver.py
+++ b/samcli/lib/intrinsic_resolver/intrinsic_property_resolver.py
@@ -492,7 +492,7 @@ def handle_find_in_map(self, intrinsic_value, ignore_errors):
         )
 
         second_level_value = top_level_value.get(second_level_key)
-        verify_intrinsic_type_str(
+        verify_non_null(
             second_level_value,
             IntrinsicResolver.FN_FIND_IN_MAP,
             message="The SecondLevelKey is missing in the Mappings dictionary in Fn::FindInMap  "
diff --git a/tests/integration/local/start_lambda/test_start_lambda.py b/tests/integration/local/start_lambda/test_start_lambda.py
index 746211e4b4..f7243bdda8 100644
--- a/tests/integration/local/start_lambda/test_start_lambda.py
+++ b/tests/integration/local/start_lambda/test_start_lambda.py
@@ -286,6 +286,28 @@ def test_invoke_with_function_timeout(self, use_full_path):
         self.assertIsNone(response.get("FunctionError"))
         self.assertEqual(response.get("StatusCode"), 200)
 
+    @parameterized.expand([("False"), ("True")])
+    @pytest.mark.flaky(reruns=3)
+    @pytest.mark.timeout(timeout=300, method="thread")
+    def test_invoke_with_function_timeout_using_lookup_value(self, use_full_path):
+        """
+        This behavior does not match the actually Lambda Service. For functions that timeout, data returned like the
+        following:
+        {"errorMessage":"<timestamp> <request_id> Task timed out after 5.00 seconds"}
+
+        For Local Lambda's, however, timeouts are an interrupt on the thread that runs invokes the function. Since the
+        invoke is on a different thread, we do not (currently) have a way to communicate this back to the caller. So
+        when a timeout happens locally, we do not add the FunctionError: Unhandled to the response and have an empty
+        string as the data returned (because no data was found in stdout from the container).
+        """
+        response = self.lambda_client.invoke(
+            FunctionName=f"{self.parent_path if use_full_path == 'True' else ''}TimeoutFunctionUsingLookupValue"
+        )
+
+        self.assertEqual(response.get("Payload").read().decode("utf-8"), "")
+        self.assertIsNone(response.get("FunctionError"))
+        self.assertEqual(response.get("StatusCode"), 200)
+
 
 class TestWarmContainersBaseClass(StartLambdaIntegBaseClass):
     def setUp(self):
diff --git a/tests/integration/testdata/invoke/template.yml b/tests/integration/testdata/invoke/template.yml
index 30ab827a5d..76db225f5c 100644
--- a/tests/integration/testdata/invoke/template.yml
+++ b/tests/integration/testdata/invoke/template.yml
@@ -26,6 +26,11 @@ Parameters:
     Type: String
     Default: "2"
 
+Mappings:
+  common:
+    LambdaFunction:
+      Timeout: 5
+
 Resources:
   HelloWorldServerlessFunction:
     Type: AWS::Serverless::Function
@@ -58,6 +63,14 @@ Resources:
       CodeUri: .
       Timeout: 5
 
+  TimeoutFunctionUsingLookupValue:
+    Type: AWS::Serverless::Function
+    Properties:
+      Handler: main.sleep_handler
+      Runtime: python3.9
+      CodeUri: .
+      Timeout: !FindInMap [common, LambdaFunction, Timeout]    
+
   HelloWorldSleepFunction:
     Type: AWS::Serverless::Function
     Properties:
diff --git a/tests/unit/lib/intrinsic_resolver/test_intrinsic_resolver.py b/tests/unit/lib/intrinsic_resolver/test_intrinsic_resolver.py
index e7d4be3b5f..22e58d7875 100644
--- a/tests/unit/lib/intrinsic_resolver/test_intrinsic_resolver.py
+++ b/tests/unit/lib/intrinsic_resolver/test_intrinsic_resolver.py
@@ -209,6 +209,7 @@ def setUp(self):
                 "Basic": {"Test": {"key": "value"}},
                 "value": {"anotherkey": {"key": "result"}},
                 "result": {"value": {"key": "final"}},
+                "NonStrValue": {"Test": {"key": 0}},
             }
         }
         self.resolver = IntrinsicResolver(symbol_resolver=IntrinsicsSymbolTable(), template=template)
@@ -218,6 +219,11 @@ def test_basic_find_in_map(self):
         result = self.resolver.intrinsic_property_resolver(intrinsic, True)
         self.assertEqual(result, "value")
 
+    def test_basic_find_in_map_with_non_string_value(self):
+        intrinsic = {"Fn::FindInMap": ["NonStrValue", "Test", "key"]}
+        result = self.resolver.intrinsic_property_resolver(intrinsic, True)
+        self.assertEqual(result, 0)
+
     def test_nested_find_in_map(self):
         intrinsic_base_1 = {"Fn::FindInMap": ["Basic", "Test", "key"]}
         intrinsic_base_2 = {"Fn::FindInMap": [intrinsic_base_1, "anotherkey", "key"]}