amoghrajesh commented on code in PR #58638:
URL: https://github.com/apache/airflow/pull/58638#discussion_r2565177109
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -174,61 +192,246 @@ def test_xcom_delete(sdk_client, dag_info):
console.print("[green]✅ XCom delete test passed!")
[email protected](reason="TODO: Implement XCom head test")
def test_xcom_head(sdk_client, dag_info):
"""
Test getting count of mapped XCom values.
Expected: XComCountResponse with len field in it (should be ideally equal
to number of mapped tasks, since we have None it might throw RuntimeError)
Endpoint: HEAD /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_xcom_head")
- raise NotImplementedError("test_xcom_head not implemented")
+ console.print("[yellow]Testing XCom head for non-mapped task...")
+
+ response_single = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="return_tuple_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Non-Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_single).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_single.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_single, XComCountResponse)
+ assert response_single.len == 1
+
+ console.print("[yellow]Testing XCom head for mapped task...")
+
+ response_mapped = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_mapped).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_mapped.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_mapped, XComCountResponse)
+ assert response_mapped.len == 4
+ console.print("[green]✅ XCom head test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item test")
def test_xcom_get_sequence_item(sdk_client, dag_info):
"""
Test getting XCom sequence item by offset.
Expected: XComSequenceIndexResponse with value
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement test_xcom_get_sequence_item")
- raise NotImplementedError("test_xcom_get_sequence_item not implemented")
+ console.print("[yellow]Testing XCom sequence item access...")
+
+ response_0 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item [offset=0] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_0).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_0.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_0, XComSequenceIndexResponse)
+ assert response_0.root == "processed_alpha"
+
+ response_neg1 = sdk_client.xcoms.get_sequence_item(
Review Comment:
Why `neg`?
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -174,61 +192,246 @@ def test_xcom_delete(sdk_client, dag_info):
console.print("[green]✅ XCom delete test passed!")
[email protected](reason="TODO: Implement XCom head test")
def test_xcom_head(sdk_client, dag_info):
"""
Test getting count of mapped XCom values.
Expected: XComCountResponse with len field in it (should be ideally equal
to number of mapped tasks, since we have None it might throw RuntimeError)
Endpoint: HEAD /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_xcom_head")
- raise NotImplementedError("test_xcom_head not implemented")
+ console.print("[yellow]Testing XCom head for non-mapped task...")
+
+ response_single = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="return_tuple_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Non-Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_single).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_single.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_single, XComCountResponse)
+ assert response_single.len == 1
+
+ console.print("[yellow]Testing XCom head for mapped task...")
+
+ response_mapped = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_mapped).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_mapped.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_mapped, XComCountResponse)
+ assert response_mapped.len == 4
+ console.print("[green]✅ XCom head test passed!")
Review Comment:
Lets split into two tests
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -174,61 +192,246 @@ def test_xcom_delete(sdk_client, dag_info):
console.print("[green]✅ XCom delete test passed!")
[email protected](reason="TODO: Implement XCom head test")
def test_xcom_head(sdk_client, dag_info):
"""
Test getting count of mapped XCom values.
Expected: XComCountResponse with len field in it (should be ideally equal
to number of mapped tasks, since we have None it might throw RuntimeError)
Endpoint: HEAD /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_xcom_head")
- raise NotImplementedError("test_xcom_head not implemented")
+ console.print("[yellow]Testing XCom head for non-mapped task...")
+
+ response_single = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="return_tuple_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Non-Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_single).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_single.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_single, XComCountResponse)
+ assert response_single.len == 1
+
+ console.print("[yellow]Testing XCom head for mapped task...")
+
+ response_mapped = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_mapped).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_mapped.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_mapped, XComCountResponse)
+ assert response_mapped.len == 4
+ console.print("[green]✅ XCom head test passed!")
Review Comment:
One for mapped and one for unmapped
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -174,61 +192,246 @@ def test_xcom_delete(sdk_client, dag_info):
console.print("[green]✅ XCom delete test passed!")
[email protected](reason="TODO: Implement XCom head test")
def test_xcom_head(sdk_client, dag_info):
"""
Test getting count of mapped XCom values.
Expected: XComCountResponse with len field in it (should be ideally equal
to number of mapped tasks, since we have None it might throw RuntimeError)
Endpoint: HEAD /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_xcom_head")
- raise NotImplementedError("test_xcom_head not implemented")
+ console.print("[yellow]Testing XCom head for non-mapped task...")
+
+ response_single = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="return_tuple_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Non-Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_single).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_single.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_single, XComCountResponse)
+ assert response_single.len == 1
+
+ console.print("[yellow]Testing XCom head for mapped task...")
+
+ response_mapped = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_mapped).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_mapped.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_mapped, XComCountResponse)
+ assert response_mapped.len == 4
+ console.print("[green]✅ XCom head test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item test")
def test_xcom_get_sequence_item(sdk_client, dag_info):
"""
Test getting XCom sequence item by offset.
Expected: XComSequenceIndexResponse with value
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement test_xcom_get_sequence_item")
- raise NotImplementedError("test_xcom_get_sequence_item not implemented")
+ console.print("[yellow]Testing XCom sequence item access...")
+
+ response_0 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item [offset=0] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_0).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_0.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_0, XComSequenceIndexResponse)
+ assert response_0.root == "processed_alpha"
+
+ response_neg1 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=-1,
+ )
+
+ console.print(" XCom Sequence Item [offset=-1] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_neg1).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_neg1.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_neg1, XComSequenceIndexResponse)
+ assert response_neg1.root == "processed_delta"
+
+ response_2 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=2,
+ )
+
+ console.print(" XCom Sequence Item [offset=2] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_2).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_2.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_2, XComSequenceIndexResponse)
+ assert response_2.root == "processed_gamma" # Third mapped instance
(index 2)
+
+ console.print("[green]✅ XCom get_sequence_item test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item (not found)
test")
def test_xcom_get_sequence_item_not_found(sdk_client, dag_info):
"""
Test getting non-existent XCom sequence item.
Expected: ErrorResponse with XCOM_NOT_FOUND error
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement
test_xcom_get_sequence_item_not_found")
- raise NotImplementedError("test_xcom_get_sequence_item_not_found not
implemented")
+ console.print("[yellow]Testing XCom sequence item not found...")
+
+ response = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=10,
+ )
+
+ console.print(" XCom Sequence Item Not Found Response ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/] {type(response).__name__}")
+ console.print(f"[bright_blue]Error:[/] {response.error}")
+ console.print(f"[bright_blue]Detail:[/] {response.detail}")
+ console.print("=" * 72)
+
+ assert isinstance(response, ErrorResponse)
+ assert response.error == ErrorType.XCOM_NOT_FOUND
+ assert response.detail["key"] == "return_value"
+ assert response.detail["offset"] == 10
+
+ response_bad_key = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="non_existent_key",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item (Bad Key) Response ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_bad_key).__name__}")
+ console.print(f"[bright_blue]Error:[/] {response_bad_key.error}")
+ console.print("=" * 72)
+
+ assert isinstance(response_bad_key, ErrorResponse)
+ assert response_bad_key.error == ErrorType.XCOM_NOT_FOUND
+
+ console.print("[green]✅ XCom get_sequence_item_not_found test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_slice test")
def test_xcom_get_sequence_slice(sdk_client, dag_info):
"""
Test getting XCom sequence slice.
Expected: XComSequenceSliceResponse with list of values
Endpoint: GET /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/slice
"""
- console.print("[yellow]TODO: Implement test_xcom_get_sequence_slice")
- raise NotImplementedError("test_xcom_get_sequence_slice not implemented")
+ console.print("[yellow]Testing XCom sequence slice access...")
+
+ response_full = sdk_client.xcoms.get_sequence_slice(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ start=None,
+ stop=None,
+ step=None,
+ )
+
+ console.print(" XCom Sequence Slice (full) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_full).__name__}")
+ console.print(f"[bright_blue]Values:[/] {response_full.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_full, XComSequenceSliceResponse)
+ assert response_full.root == ["processed_alpha", "processed_beta",
"processed_gamma", "processed_delta"]
+
+ response_slice = sdk_client.xcoms.get_sequence_slice(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ start=1,
+ stop=3,
+ step=None,
+ )
+
+ console.print(" XCom Sequence Slice [start=1,stop=3] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_slice).__name__}")
+ console.print(f"[bright_blue]Values:[/] {response_slice.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_slice, XComSequenceSliceResponse)
+ assert response_slice.root == ["processed_beta", "processed_gamma"]
+
+ response_step = sdk_client.xcoms.get_sequence_slice(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ start=0,
+ stop=4,
+ step=2,
+ )
+
+ console.print(" XCom Sequence Slice [step=2] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_step).__name__}")
+ console.print(f"[bright_blue]Values:[/] {response_step.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_step, XComSequenceSliceResponse)
+ assert response_step.root == ["processed_alpha", "processed_gamma"]
+
+ console.print("[green]✅ XCom get_sequence_slice test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_slice (not found)
test")
def test_xcom_get_sequence_slice_not_found(sdk_client, dag_info):
"""
Test getting slice for non-existent XCom key.
Expected: XComSequenceSliceResponse as empty list
Endpoint: GET /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/slice
"""
- console.print("[yellow]TODO: Implement
test_xcom_get_sequence_slice_not_found")
- raise NotImplementedError("test_xcom_get_sequence_slice_not_found not
implemented")
+ console.print("[yellow]Testing XCom sequence slice not found...")
+
+ response = sdk_client.xcoms.get_sequence_slice(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="non_existent_key",
+ start=0,
+ stop=10,
+ step=None,
+ )
+
+ console.print(" XCom Sequence Slice (Not Found) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/] {type(response).__name__}")
+ console.print(f"[bright_blue]Values:[/] {getattr(response, 'root', None)}")
+ console.print("=" * 72)
+
+ assert isinstance(response, XComSequenceSliceResponse)
+ assert response.root == []
+ console.print("[green]✅ XCom get_sequence_slice_not_found test passed!")
Review Comment:
This looks good.
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -174,61 +192,246 @@ def test_xcom_delete(sdk_client, dag_info):
console.print("[green]✅ XCom delete test passed!")
[email protected](reason="TODO: Implement XCom head test")
def test_xcom_head(sdk_client, dag_info):
"""
Test getting count of mapped XCom values.
Expected: XComCountResponse with len field in it (should be ideally equal
to number of mapped tasks, since we have None it might throw RuntimeError)
Endpoint: HEAD /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_xcom_head")
- raise NotImplementedError("test_xcom_head not implemented")
+ console.print("[yellow]Testing XCom head for non-mapped task...")
+
+ response_single = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="return_tuple_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Non-Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_single).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_single.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_single, XComCountResponse)
+ assert response_single.len == 1
+
+ console.print("[yellow]Testing XCom head for mapped task...")
+
+ response_mapped = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_mapped).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_mapped.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_mapped, XComCountResponse)
+ assert response_mapped.len == 4
+ console.print("[green]✅ XCom head test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item test")
def test_xcom_get_sequence_item(sdk_client, dag_info):
"""
Test getting XCom sequence item by offset.
Expected: XComSequenceIndexResponse with value
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement test_xcom_get_sequence_item")
- raise NotImplementedError("test_xcom_get_sequence_item not implemented")
+ console.print("[yellow]Testing XCom sequence item access...")
+
+ response_0 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item [offset=0] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_0).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_0.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_0, XComSequenceIndexResponse)
+ assert response_0.root == "processed_alpha"
+
+ response_neg1 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=-1,
+ )
+
+ console.print(" XCom Sequence Item [offset=-1] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_neg1).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_neg1.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_neg1, XComSequenceIndexResponse)
+ assert response_neg1.root == "processed_delta"
+
+ response_2 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=2,
+ )
+
+ console.print(" XCom Sequence Item [offset=2] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_2).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_2.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_2, XComSequenceIndexResponse)
+ assert response_2.root == "processed_gamma" # Third mapped instance
(index 2)
+
+ console.print("[green]✅ XCom get_sequence_item test passed!")
Review Comment:
Guess we can just run this thing in a loop?
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -64,16 +67,33 @@ def test_get_xcom(sdk_client, dag_info):
console.print("[green]✅ XCom get test passed!")
[email protected](reason="TODO: Implement XCom get (not found) test")
def test_get_xcom_not_found(sdk_client, dag_info):
"""
Test getting non-existent XCom value.
Expected: XComResponse with value=None or ErrorResponse
Endpoint: GET /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_get_xcom_not_found")
- raise NotImplementedError("test_get_xcom_not_found not implemented")
+ missing_key = "non_existent_xcom_key_for_test"
+ console.print("[yellow]Getting non-existent XCom key...")
+
+ response = sdk_client.xcoms.get(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="get_task_instance_id",
+ key=missing_key,
+ )
+
+ console.print(" XCom Get (Not Found) Response ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/] {type(response).__name__}")
+ console.print(f"[bright_blue]Key:[/] {response.key}")
+ console.print(f"[bright_blue]Value:[/] {response.value}")
+ console.print("=" * 72)
+
+ assert isinstance(response, XComResponse)
+ assert response.key == missing_key
+ assert response.value is None
+ console.print("[green]✅ XCom not-found test passed!")
Review Comment:
This looks good.
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -174,61 +192,246 @@ def test_xcom_delete(sdk_client, dag_info):
console.print("[green]✅ XCom delete test passed!")
[email protected](reason="TODO: Implement XCom head test")
def test_xcom_head(sdk_client, dag_info):
"""
Test getting count of mapped XCom values.
Expected: XComCountResponse with len field in it (should be ideally equal
to number of mapped tasks, since we have None it might throw RuntimeError)
Endpoint: HEAD /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_xcom_head")
- raise NotImplementedError("test_xcom_head not implemented")
+ console.print("[yellow]Testing XCom head for non-mapped task...")
+
+ response_single = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="return_tuple_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Non-Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_single).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_single.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_single, XComCountResponse)
+ assert response_single.len == 1
+
+ console.print("[yellow]Testing XCom head for mapped task...")
+
+ response_mapped = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_mapped).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_mapped.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_mapped, XComCountResponse)
+ assert response_mapped.len == 4
+ console.print("[green]✅ XCom head test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item test")
def test_xcom_get_sequence_item(sdk_client, dag_info):
"""
Test getting XCom sequence item by offset.
Expected: XComSequenceIndexResponse with value
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement test_xcom_get_sequence_item")
- raise NotImplementedError("test_xcom_get_sequence_item not implemented")
+ console.print("[yellow]Testing XCom sequence item access...")
+
+ response_0 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item [offset=0] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_0).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_0.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_0, XComSequenceIndexResponse)
+ assert response_0.root == "processed_alpha"
+
+ response_neg1 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=-1,
+ )
+
+ console.print(" XCom Sequence Item [offset=-1] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_neg1).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_neg1.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_neg1, XComSequenceIndexResponse)
+ assert response_neg1.root == "processed_delta"
+
+ response_2 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=2,
+ )
+
+ console.print(" XCom Sequence Item [offset=2] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_2).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_2.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_2, XComSequenceIndexResponse)
+ assert response_2.root == "processed_gamma" # Third mapped instance
(index 2)
+
+ console.print("[green]✅ XCom get_sequence_item test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item (not found)
test")
def test_xcom_get_sequence_item_not_found(sdk_client, dag_info):
"""
Test getting non-existent XCom sequence item.
Expected: ErrorResponse with XCOM_NOT_FOUND error
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement
test_xcom_get_sequence_item_not_found")
- raise NotImplementedError("test_xcom_get_sequence_item_not_found not
implemented")
+ console.print("[yellow]Testing XCom sequence item not found...")
+
+ response = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=10,
+ )
+
+ console.print(" XCom Sequence Item Not Found Response ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/] {type(response).__name__}")
+ console.print(f"[bright_blue]Error:[/] {response.error}")
+ console.print(f"[bright_blue]Detail:[/] {response.detail}")
+ console.print("=" * 72)
+
+ assert isinstance(response, ErrorResponse)
+ assert response.error == ErrorType.XCOM_NOT_FOUND
+ assert response.detail["key"] == "return_value"
+ assert response.detail["offset"] == 10
+
+ response_bad_key = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="non_existent_key",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item (Bad Key) Response ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_bad_key).__name__}")
+ console.print(f"[bright_blue]Error:[/] {response_bad_key.error}")
+ console.print("=" * 72)
+
+ assert isinstance(response_bad_key, ErrorResponse)
+ assert response_bad_key.error == ErrorType.XCOM_NOT_FOUND
+
+ console.print("[green]✅ XCom get_sequence_item_not_found test passed!")
Review Comment:
I think we should split this one into two too. One for 10 offset and one of
wrong key
##########
task-sdk-integration-tests/tests/task_sdk_tests/test_xcom_operations.py:
##########
@@ -174,61 +192,246 @@ def test_xcom_delete(sdk_client, dag_info):
console.print("[green]✅ XCom delete test passed!")
[email protected](reason="TODO: Implement XCom head test")
def test_xcom_head(sdk_client, dag_info):
"""
Test getting count of mapped XCom values.
Expected: XComCountResponse with len field in it (should be ideally equal
to number of mapped tasks, since we have None it might throw RuntimeError)
Endpoint: HEAD /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}
"""
- console.print("[yellow]TODO: Implement test_xcom_head")
- raise NotImplementedError("test_xcom_head not implemented")
+ console.print("[yellow]Testing XCom head for non-mapped task...")
+
+ response_single = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="return_tuple_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Non-Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_single).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_single.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_single, XComCountResponse)
+ assert response_single.len == 1
+
+ console.print("[yellow]Testing XCom head for mapped task...")
+
+ response_mapped = sdk_client.xcoms.head(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ )
+
+ console.print(" XCom Head Response (Mapped) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_mapped).__name__}")
+ console.print(f"[bright_blue]Count:[/] {response_mapped.len}")
+ console.print("=" * 72)
+
+ assert isinstance(response_mapped, XComCountResponse)
+ assert response_mapped.len == 4
+ console.print("[green]✅ XCom head test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item test")
def test_xcom_get_sequence_item(sdk_client, dag_info):
"""
Test getting XCom sequence item by offset.
Expected: XComSequenceIndexResponse with value
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement test_xcom_get_sequence_item")
- raise NotImplementedError("test_xcom_get_sequence_item not implemented")
+ console.print("[yellow]Testing XCom sequence item access...")
+
+ response_0 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item [offset=0] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_0).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_0.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_0, XComSequenceIndexResponse)
+ assert response_0.root == "processed_alpha"
+
+ response_neg1 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=-1,
+ )
+
+ console.print(" XCom Sequence Item [offset=-1] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_neg1).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_neg1.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_neg1, XComSequenceIndexResponse)
+ assert response_neg1.root == "processed_delta"
+
+ response_2 = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=2,
+ )
+
+ console.print(" XCom Sequence Item [offset=2] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_2).__name__}")
+ console.print(f"[bright_blue]Value:[/] {response_2.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_2, XComSequenceIndexResponse)
+ assert response_2.root == "processed_gamma" # Third mapped instance
(index 2)
+
+ console.print("[green]✅ XCom get_sequence_item test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_item (not found)
test")
def test_xcom_get_sequence_item_not_found(sdk_client, dag_info):
"""
Test getting non-existent XCom sequence item.
Expected: ErrorResponse with XCOM_NOT_FOUND error
Endpoint: GET
/execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}
"""
- console.print("[yellow]TODO: Implement
test_xcom_get_sequence_item_not_found")
- raise NotImplementedError("test_xcom_get_sequence_item_not_found not
implemented")
+ console.print("[yellow]Testing XCom sequence item not found...")
+
+ response = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ offset=10,
+ )
+
+ console.print(" XCom Sequence Item Not Found Response ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/] {type(response).__name__}")
+ console.print(f"[bright_blue]Error:[/] {response.error}")
+ console.print(f"[bright_blue]Detail:[/] {response.detail}")
+ console.print("=" * 72)
+
+ assert isinstance(response, ErrorResponse)
+ assert response.error == ErrorType.XCOM_NOT_FOUND
+ assert response.detail["key"] == "return_value"
+ assert response.detail["offset"] == 10
+
+ response_bad_key = sdk_client.xcoms.get_sequence_item(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="non_existent_key",
+ offset=0,
+ )
+
+ console.print(" XCom Sequence Item (Bad Key) Response ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_bad_key).__name__}")
+ console.print(f"[bright_blue]Error:[/] {response_bad_key.error}")
+ console.print("=" * 72)
+
+ assert isinstance(response_bad_key, ErrorResponse)
+ assert response_bad_key.error == ErrorType.XCOM_NOT_FOUND
+
+ console.print("[green]✅ XCom get_sequence_item_not_found test passed!")
[email protected](reason="TODO: Implement XCom get_sequence_slice test")
def test_xcom_get_sequence_slice(sdk_client, dag_info):
"""
Test getting XCom sequence slice.
Expected: XComSequenceSliceResponse with list of values
Endpoint: GET /execution/xcoms/{dag_id}/{run_id}/{task_id}/{key}/slice
"""
- console.print("[yellow]TODO: Implement test_xcom_get_sequence_slice")
- raise NotImplementedError("test_xcom_get_sequence_slice not implemented")
+ console.print("[yellow]Testing XCom sequence slice access...")
+
+ response_full = sdk_client.xcoms.get_sequence_slice(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ start=None,
+ stop=None,
+ step=None,
+ )
+
+ console.print(" XCom Sequence Slice (full) ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_full).__name__}")
+ console.print(f"[bright_blue]Values:[/] {response_full.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_full, XComSequenceSliceResponse)
+ assert response_full.root == ["processed_alpha", "processed_beta",
"processed_gamma", "processed_delta"]
+
+ response_slice = sdk_client.xcoms.get_sequence_slice(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ start=1,
+ stop=3,
+ step=None,
+ )
+
+ console.print(" XCom Sequence Slice [start=1,stop=3] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_slice).__name__}")
+ console.print(f"[bright_blue]Values:[/] {response_slice.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_slice, XComSequenceSliceResponse)
+ assert response_slice.root == ["processed_beta", "processed_gamma"]
+
+ response_step = sdk_client.xcoms.get_sequence_slice(
+ dag_id=dag_info["dag_id"],
+ run_id=dag_info["dag_run_id"],
+ task_id="mapped_task",
+ key="return_value",
+ start=0,
+ stop=4,
+ step=2,
+ )
+
+ console.print(" XCom Sequence Slice [step=2] ".center(72, "="))
+ console.print(f"[bright_blue]Response Type:[/]
{type(response_step).__name__}")
+ console.print(f"[bright_blue]Values:[/] {response_step.root}")
+ console.print("=" * 72)
+
+ assert isinstance(response_step, XComSequenceSliceResponse)
+ assert response_step.root == ["processed_alpha", "processed_gamma"]
+
Review Comment:
Same for this test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
