🐛fix: Only replace implementation with dynamic

terraform/modules/content_delivery_network/lambda_functions/redirect_s3_404/src/lambda_function.py

@@ -1,10 +1,13 @@
 # mypy: ignore-errors
+from urllib.parse import urlparse, urlunparse
 
 
 def handler(event, context):
-    """
-    This function updates the HTTP status code in the response to 307, to redirect to another
-    path (cache behavior) that has a different origin configured. Note the following:
+    """This function updates the HTTP status code in the response to 307, to
+    redirect to another path (cache behavior) that has a different origin
+    configured.
+
+    Note the following:
     1. The function is triggered in an origin response
     2. The response status from the origin server is an error status code (404)
     """
@@ -17,20 +20,26 @@ def handler(event, context):
     # custom origin is tile cache app. URL is passed via custom header set in cloud front
     # (env variables are not support for Lambda@Edge)
 
-    if int(response["status"]) == 404 and is_tile(request["uri"]):
+    parsed_url = urlparse(request["uri"])
+    path_parts = parsed_url.path.split("/")
 
-        implementation = get_implementation(request["uri"])
+    if int(response["status"]) == 404 and is_tile(path_parts):
+        implementation = replace_implementation(path_parts)
 
-        redirect_path = request["uri"].replace(implementation, "dynamic")
+        updated_path = "/".join(path_parts)
 
+        # Append the query string with the `implementation`
         if request["querystring"]:
             querystring = f"{request['querystring']}&implementation={implementation}"
         else:
             querystring = f"implementation={implementation}"
 
-        redirect_path += f"?{querystring}"
+        updated_url = urlunparse(
+            parsed_url._replace(path=updated_path, query=querystring)
+        )
 
-        headers["location"] = [{"key": "Location", "value": redirect_path}]
+        # Set the headers for the redirect response
+        headers["location"] = [{"key": "Location", "value": updated_url}]
         headers["content-type"] = [{"key": "Content-Type", "value": "application/json"}]
         headers["content-encoding"] = [{"key": "Content-Encoding", "value": "UTF-8"}]
 
@@ -44,11 +53,51 @@ def handler(event, context):
 
 
 def is_tile(uri):
-    print("REQUEST URI", uri)
-    parts = uri.split("/")
-    return len(parts) == 7 and parts[6][-4:] in [".png", ".pbf"]
-
-
-def get_implementation(uri):
-    parts = uri.split("/")
-    return parts[3]
+    print("REQUEST URI", "/".join(uri))
+    return len(uri) == 7 and uri[6][-4:] in [".png", ".pbf"]
+
+
+def replace_implementation(path_parts):
+    """Replace the implementation path segment with "dynamic" and return the
+    original implementation."""
+    impl_path_location = 3
+    implementation = path_parts[impl_path_location]
+    path_parts[impl_path_location] = "dynamic"
+    return implementation
+
+
+def test_handler():
+    # Mock event and context
+    event = {
+        "Records": [
+            {
+                "cf": {
+                    "response": {
+                        "status": "404",
+                        "headers": {
+                            "content-type": [
+                                {"key": "Content-Type", "value": "text/plain"}
+                            ]
+                        },
+                    },
+                    "request": {
+                        "uri": "/sbtn_natural_forests_map/v202310/natural_forest/10/20/30.png",
+                        "querystring": "some_param=30",
+                    },
+                }
+            }
+        ]
+    }
+    context = {}
+
+    # Call the handler function
+    response = handler(event, context)
+
+    # Assertions
+    assert response["status"] == "307"
+    assert response["headers"]["location"][0]["key"] == "Location"
+    assert (
+        response["headers"]["location"][0]["value"]
+        == "/sbtn_natural_forests_map/v202310/dynamic/10/20/30.png?some_param=30&implementation=natural_forest"
+    )
+    assert response["headers"]["content-type"][0]["value"] == "application/json"