From 1b6a648a5cdab47364d918b434947c381787350e Mon Sep 17 00:00:00 2001
From: Brent Champion <brchamp@amazon.com>
Date: Wed, 8 Oct 2025 13:57:43 -0400
Subject: [PATCH] fix: update exception handling in PutLambdaEndpoint API

---
 examples/cli.py                                        |  6 ++++--
 .../web/handlers.py                                    | 10 ++++------
 tests/web/handlers_test.py                             |  3 ++-
 3 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/examples/cli.py b/examples/cli.py
index c501531..fedf2a0 100755
--- a/examples/cli.py
+++ b/examples/cli.py
@@ -226,7 +226,7 @@ def generate_sam_template():
                 "Timeout": 60,
                 "MemorySize": 128,
                 "Environment": {
-                    "Variables": {"DEX_ENDPOINT": {"Ref": "LambdaEndpoint"}}
+                    "Variables": {"AWS_ENDPOINT_URL_LAMBDA": {"Ref": "LambdaEndpoint"}}
                 },
             }
         },
@@ -353,7 +353,9 @@ def deploy_function(example_name: str, function_name: str | None = None):
         "Description": example_config["description"],
         "Timeout": 60,
         "MemorySize": 128,
-        # "Environment": {"Variables": {"AWS_ENDPOINT_URL_LAMBDA": config["lambda_endpoint"]}},
+        "Environment": {
+            "Variables": {"AWS_ENDPOINT_URL_LAMBDA": config["lambda_endpoint"]}
+        },
         "DurableConfig": example_config["durableConfig"],
     }
 
diff --git a/src/aws_durable_execution_sdk_python_testing/web/handlers.py b/src/aws_durable_execution_sdk_python_testing/web/handlers.py
index 6eb395b..39d1df8 100644
--- a/src/aws_durable_execution_sdk_python_testing/web/handlers.py
+++ b/src/aws_durable_execution_sdk_python_testing/web/handlers.py
@@ -790,8 +790,8 @@ def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse:  #
             region_name = body.get("RegionName", "us-east-1")
 
             if not endpoint_url:
-                return HTTPResponse.create_json(
-                    400, {"error": "EndpointUrl is required"}
+                return self._handle_aws_exception(
+                    InvalidParameterValueException("EndpointUrl is required")
                 )
 
             # Update the invoker's Lambda endpoint
@@ -802,7 +802,5 @@ def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse:  #
                 {"message": "Lambda endpoint updated successfully"}
             )
 
-        except (AttributeError, TypeError) as e:
-            return HTTPResponse.create_json(
-                500, {"error": f"Failed to update Lambda endpoint: {e!s}"}
-            )
+        except Exception as e:  # noqa: BLE001
+            return self._handle_framework_exception(e)
diff --git a/tests/web/handlers_test.py b/tests/web/handlers_test.py
index dfd4b91..40bd966 100644
--- a/tests/web/handlers_test.py
+++ b/tests/web/handlers_test.py
@@ -2109,7 +2109,8 @@ def test_update_lambda_endpoint_handler_missing_endpoint_url():
     response = handler.handle(update_route, request)
 
     assert response.status_code == 400
-    assert response.body == {"error": "EndpointUrl is required"}
+    assert response.body["Type"] == "InvalidParameterValueException"
+    assert response.body["message"] == "EndpointUrl is required"
 
 
 def test_update_lambda_endpoint_handler_default_region():