fix: retry transient get durable execution not found (#512)

Author: zhongkechen

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/runner.py

@@ -1054,6 +1054,17 @@ def _wait_for_completion(
                     DurableExecutionArn=execution_arn
                 )
                 execution = GetDurableExecutionResponse.from_dict(execution_dict)
+            except ClientError as e:
+                error_code = e.response.get("Error", {}).get("Code")
+                # Retry while the asynchronous invoke's durable execution record is
+                # still becoming visible to GetDurableExecution.
+                if error_code == "ResourceNotFoundException":
+                    logger.info("Execution status not available yet; retrying")
+                    time.sleep(self.poll_interval)
+                    continue
+
+                msg = f"Failed to get execution status: {e}"
+                raise DurableFunctionsTestError(msg) from e
             except Exception as e:
                 msg = f"Failed to get execution status: {e}"
                 raise DurableFunctionsTestError(msg) from e

packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py

@@ -1386,6 +1386,43 @@ def test_cloud_runner_run_missing_execution_arn(mock_boto3):
         runner.run(input="test-input")
 
 
+@patch("aws_durable_execution_sdk_python_testing.runner.boto3")
+def test_cloud_runner_wait_for_completion_get_execution_resource_not_found_retries(
+    mock_boto3,
+):
+    """Test _wait_for_completion retries transient ResourceNotFoundException."""
+    from botocore.exceptions import ClientError  # type: ignore
+
+    from aws_durable_execution_sdk_python_testing.runner import (
+        DurableFunctionCloudTestRunner,
+    )
+
+    mock_client = Mock()
+    mock_boto3.client.return_value = mock_client
+    mock_client.get_durable_execution.side_effect = [
+        ClientError(
+            error_response={"Error": {"Code": "ResourceNotFoundException"}},
+            operation_name="GetDurableExecution",
+        ),
+        {
+            "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:test:execution:exec-1",
+            "DurableExecutionName": "test-execution",
+            "FunctionArn": "arn:aws:lambda:us-east-1:123456789012:function:test",
+            "Status": "SUCCEEDED",
+            "StartTimestamp": "2023-01-01T00:00:00Z",
+            "EndTimestamp": "2023-01-01T00:01:00Z",
+        },
+    ]
+
+    runner = DurableFunctionCloudTestRunner(
+        function_name="test-function", poll_interval=0
+    )
+    result = runner._wait_for_completion("test-arn", timeout=10)
+
+    assert result.status == "SUCCEEDED"
+    assert mock_client.get_durable_execution.call_count == 2
+
+
 @patch("aws_durable_execution_sdk_python_testing.runner.boto3")
 def test_cloud_runner_wait_for_completion_get_execution_failure(mock_boto3):
     """Test DurableFunctionCloudTestRunner._wait_for_completion with API failure."""