From 2b33d5d3a4c8b49e127192ed30b854f7b6707d53 Mon Sep 17 00:00:00 2001 From: kiwigitops Date: Tue, 16 Jun 2026 19:46:12 -0400 Subject: [PATCH] fix: set wait_for_callback context subtype --- .../src/aws_durable_execution_sdk_python/context.py | 1 + .../aws-durable-execution-sdk-python/tests/context_test.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py index 2f9bc21e..5f14f605 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py @@ -717,6 +717,7 @@ def wait_in_child_context(context: DurableContext): return self.run_in_child_context( wait_in_child_context, step_name, + ChildConfig(sub_type=OperationSubType.WAIT_FOR_CALLBACK), ) def wait_for_condition( diff --git a/packages/aws-durable-execution-sdk-python/tests/context_test.py b/packages/aws-durable-execution-sdk-python/tests/context_test.py index 13d1dc64..42c0205b 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -1196,6 +1196,7 @@ def test_wait_for_callback_basic(mock_executor_class): # Verify the child context callable call_args = mock_run_in_child.call_args assert call_args[0][1] is None # name should be None + assert call_args[0][2].sub_type is OperationSubType.WAIT_FOR_CALLBACK @patch("aws_durable_execution_sdk_python.context.wait_for_callback_handler") @@ -1225,6 +1226,7 @@ def test_wait_for_callback_with_name_and_config(mock_executor_class): assert ( call_args[0][1] == "submit_function" ) # name should be from _original_name + assert call_args[0][2].sub_type is OperationSubType.WAIT_FOR_CALLBACK @patch("aws_durable_execution_sdk_python.context.wait_for_callback_handler") @@ -1250,6 +1252,7 @@ def test_wait_for_callback_resolves_name_from_submitter(mock_executor_class): call_args = mock_run_in_child.call_args assert call_args[0][1] == "submit_task" + assert call_args[0][2].sub_type is OperationSubType.WAIT_FOR_CALLBACK @patch("aws_durable_execution_sdk_python.context.wait_for_callback_handler") @@ -1270,8 +1273,9 @@ def capture_handler_call(context, submitter, name, config): with patch.object(DurableContext, "run_in_child_context") as mock_run_in_child: - def run_child_context(callable_func, name): + def run_child_context(callable_func, name, config): # Execute the child context callable + assert config.sub_type is OperationSubType.WAIT_FOR_CALLBACK child_context = create_test_context(state=mock_state, parent_id="test") return callable_func(child_context)