fix(tests): stabilize flaky durable timeout tests

.github/workflows/build.yml

@@ -197,6 +197,7 @@ jobs:
         run: pytest -vv ${{ matrix.tests_config.params }}
         env:
           FORCE_RUN_DOCKER_TEST: ${{ matrix.tests_config.name == 'durable-functions' && '1' || '' }}
+          DURABLE_EXECUTIONS_EMULATOR_IMAGE_TAG: ${{ matrix.tests_config.name == 'durable-functions' && 'v1.1.1' || '' }}
 
   smoke-and-functional-tests:
     name: ${{ matrix.tests_config.name }} / ${{ matrix.tests_config.os }} / ${{ matrix.python }}

tests/integration/local/invoke/test_invoke_durable.py

@@ -70,11 +70,8 @@ def test_local_invoke_durable_function_timeout(self):
         example = DurableFunctionExamples.EXECUTION_TIMEOUT
         function_name = example.function_name
         execution_name = "executiontimeout-integration-test"
-        event_path = str(self.test_data_path / "durable" / "events" / "timeout_test_event.json")
 
-        command_list = self.get_invoke_command_list(
-            function_name, event_path=event_path, durable_execution_name=execution_name
-        )
+        command_list = self.get_invoke_command_list(function_name, no_event=True, durable_execution_name=execution_name)
 
         stdout, stderr, invoke_return_code = self.run_command_with_logging(
             command_list, f"test_local_invoke_durable_function_{function_name}"
@@ -84,7 +81,7 @@ def test_local_invoke_durable_function_timeout(self):
 
         # Assert invoke output shows timeout
         execution_arn = self.assert_invoke_output(
-            stdout, input_data={"wait_seconds": 30}, execution_name=execution_name, expected_status="TIMED_OUT"
+            stdout, input_data={}, execution_name=execution_name, expected_status="TIMED_OUT"
         )
 
         # Get and verify execution history

tests/integration/local/start_lambda/test_start_lambda_durable.py

@@ -103,18 +103,16 @@ def wait_for_execution_status(self, execution_arn, expected_status, max_wait=30)
             time.sleep(1)
         return execution_response
 
-    def invoke_and_wait_for_callback(self, payload=None):
-        """Helper to invoke WaitForCallback function and wait for callback to be pending.
+    def invoke_and_wait_for_callback(self, payload=None, function_name="WaitForCallback"):
+        """Helper to invoke a wait_for_callback function and wait for callback to be pending.
 
         Returns:
             tuple: (execution_arn, callback_id)
         """
         if payload:
-            response = self.lambda_client.invoke(
-                FunctionName="WaitForCallback", InvocationType="Event", Payload=payload
-            )
+            response = self.lambda_client.invoke(FunctionName=function_name, InvocationType="Event", Payload=payload)
         else:
-            response = self.lambda_client.invoke(FunctionName="WaitForCallback", InvocationType="Event")
+            response = self.lambda_client.invoke(FunctionName=function_name, InvocationType="Event")
 
         execution_arn = self.assert_durable_invoke_response(
             response, DurableFunctionExamples.WAIT_FOR_CALLBACK, invocation_type="Event"
@@ -170,12 +168,10 @@ def test_local_start_lambda_invoke_timeout(self):
         """Test start-lambda with durable function execution timeout."""
         example = DurableFunctionExamples.EXECUTION_TIMEOUT
         execution_name = "executiontimeout-integration-test"
-        event_data = {"wait_seconds": 30}
 
         response = self.lambda_client.invoke(
             FunctionName=example.function_name,
             DurableExecutionName=execution_name,
-            Payload=json.dumps(event_data).encode("utf-8"),
         )
 
         self.assertEqual(response.get("StatusCode"), 200)
@@ -250,10 +246,7 @@ def test_local_start_lambda_invoke_wait_for_callback_failure_http(self, name, er
     @pytest.mark.timeout(timeout=300, method="thread")
     def test_local_start_lambda_invoke_wait_for_callback_timeout(self):
         """Test start-lambda with wait_for_callback timeout (no callback sent)."""
-        # Set a short timeout so test doesn't take too long
-        event_payload = json.dumps({"timeout_seconds": 5, "heartbeat_timeout_seconds": 3})
-
-        execution_arn, callback_id = self.invoke_and_wait_for_callback(payload=event_payload)
+        execution_arn, callback_id = self.invoke_and_wait_for_callback(function_name="WaitForCallbackShortTimeout")
 
         # Don't send any callback - let it timeout
         execution_response = self.wait_for_execution_status(execution_arn, "FAILED", max_wait=15)

tests/integration/testdata/durable/functions/wait/wait_30_seconds.py

@@ -0,0 +1,11 @@
+from typing import Any
+
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.config import Duration
+
+
+@durable_execution
+def handler(event: Any, context: DurableContext) -> str:
+    context.wait(Duration.from_seconds(30), name="custom_wait")
+    return "Wait with name completed"

tests/integration/testdata/durable/functions/wait_for_callback/wait_for_callback_short_timeout.py

@@ -0,0 +1,31 @@
+import logging
+from typing import Any
+
+from aws_durable_execution_sdk_python.config import WaitForCallbackConfig
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    WaitForCallbackContext,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.config import Duration
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+def external_system_call(callback_id: str, _context: WaitForCallbackContext) -> None:
+    logger.info(f"Waiting for callback: {callback_id}")
+
+
+@durable_execution
+def handler(event: Any, context: DurableContext) -> str:
+    config = WaitForCallbackConfig(
+        timeout=Duration.from_seconds(5),
+        heartbeat_timeout=Duration.from_seconds(3),
+    )
+
+    result = context.wait_for_callback(
+        external_system_call, name="external_call", config=config
+    )
+
+    return f"External system result: {result}"

tests/integration/testdata/durable/template.yaml

@@ -57,7 +57,13 @@ Resources:
     Type: AWS::Serverless::Function
     Properties:
       CodeUri: functions
-      Handler: wait.wait_with_name.handler
+      Handler: wait.wait_30_seconds.handler
       DurableConfig:
         ExecutionTimeout: 5
         RetentionPeriodInDays: 7
+
+  WaitForCallbackShortTimeout:
+    Type: AWS::Serverless::Function
+    Properties:
+      CodeUri: functions
+      Handler: wait_for_callback.wait_for_callback_short_timeout.handler