|
98 | 98 | _VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test" |
99 | 99 | _TASK_STATE_DAG_ID = "java_task_state" |
100 | 100 | _TASK_STATE_TASK_ID = "roundtrip_task_state" |
| 101 | +_ASSET_STATE_STORE_DAG_ID = "java_asset_state_store" |
| 102 | +_ASSET_STATE_STORE_TASK_ID = "use_asset_state_store" |
| 103 | +_ASSET_STATE_STORE_ASSET_NAME = "java_e2e_orders" |
101 | 104 |
|
102 | 105 |
|
103 | 106 | @dataclass |
@@ -226,6 +229,12 @@ def task_state_run() -> _CompletedRun: |
226 | 229 | return _trigger_and_wait_for_dag(_TASK_STATE_DAG_ID, _JAVA_TASK_TIMEOUT) |
227 | 230 |
|
228 | 231 |
|
| 232 | +@pytest.fixture(scope="module") |
| 233 | +def asset_state_store_run() -> _CompletedRun: |
| 234 | + """Trigger the asset state store Dag once for all of its assertions.""" |
| 235 | + return _trigger_and_wait_for_dag(_ASSET_STATE_STORE_DAG_ID, _JAVA_TASK_TIMEOUT) |
| 236 | + |
| 237 | + |
229 | 238 | class TestJavaSDKAnnotationExample: |
230 | 239 | """Verify the annotation-based Java SDK example executes correctly.""" |
231 | 240 |
|
@@ -427,6 +436,63 @@ def test_scratch_variable_deleted_by_java_task_is_gone(self, variable_write_run: |
427 | 436 | ) |
428 | 437 |
|
429 | 438 |
|
| 439 | +class TestJavaSDKAssetStateStore: |
| 440 | + """Verify a Java task can read and write the state store of an asset. |
| 441 | +
|
| 442 | + The task lives in the java-test-bundle fixture project (served on the |
| 443 | + dedicated "java-test" queue). It looks up the ``java_e2e_orders`` asset both |
| 444 | + by name and by URI and calls every asset state store method. Each step |
| 445 | + changes the state through one lookup and checks it through the other, so the |
| 446 | + task fails unless both lookups reach the same asset. It leaves the |
| 447 | + ``summary`` key behind for the REST API to read. |
| 448 | + """ |
| 449 | + |
| 450 | + def _assert_task_succeeded(self, run: _CompletedRun) -> None: |
| 451 | + ti = run.get_task_instance(_ASSET_STATE_STORE_TASK_ID) |
| 452 | + assert ti.get("state") == "success", ( |
| 453 | + f"Java {_ASSET_STATE_STORE_TASK_ID!r} task did not succeed.\n" |
| 454 | + f" task state : {ti.get('state')!r}\n" |
| 455 | + f" dag state : {run.state!r}\n" |
| 456 | + f" all tasks : {run.ti_states}" |
| 457 | + ) |
| 458 | + |
| 459 | + def test_value_written_by_java_task_is_readable(self, asset_state_store_run: _CompletedRun): |
| 460 | + """The value set from Java is stored on the asset and attributed to the Java task.""" |
| 461 | + self._assert_task_succeeded(asset_state_store_run) |
| 462 | + |
| 463 | + entry = asset_state_store_run.client.get_asset_state_store_entry( |
| 464 | + _ASSET_STATE_STORE_ASSET_NAME, "summary" |
| 465 | + ) |
| 466 | + assert entry["value"] == {"run_id": asset_state_store_run.run_id, "rows": 3}, ( |
| 467 | + f"summary should hold the map the Java task stored in this run, got {entry!r}" |
| 468 | + ) |
| 469 | + ti = asset_state_store_run.get_task_instance(_ASSET_STATE_STORE_TASK_ID) |
| 470 | + expected_writer = { |
| 471 | + "kind": "task", |
| 472 | + "dag_id": _ASSET_STATE_STORE_DAG_ID, |
| 473 | + "run_id": asset_state_store_run.run_id, |
| 474 | + "task_id": _ASSET_STATE_STORE_TASK_ID, |
| 475 | + "map_index": -1, |
| 476 | + "task_instance_id": ti["id"], |
| 477 | + "try_number": ti["try_number"], |
| 478 | + } |
| 479 | + writer = entry["last_updated_by"] |
| 480 | + assert {key: writer.get(key) for key in expected_writer} == expected_writer, ( |
| 481 | + f"summary should be attributed to the Java task, got {entry!r}" |
| 482 | + ) |
| 483 | + |
| 484 | + def test_key_deleted_by_java_task_is_gone(self, asset_state_store_run: _CompletedRun): |
| 485 | + """A key written and then deleted from Java no longer exists.""" |
| 486 | + self._assert_task_succeeded(asset_state_store_run) |
| 487 | + |
| 488 | + with pytest.raises(requests.HTTPError) as excinfo: |
| 489 | + asset_state_store_run.client.get_asset_state_store_entry(_ASSET_STATE_STORE_ASSET_NAME, "deleted") |
| 490 | + assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND, ( |
| 491 | + f"the 'deleted' key should have been deleted by the Java task, " |
| 492 | + f"got HTTP {excinfo.value.response.status_code}" |
| 493 | + ) |
| 494 | + |
| 495 | + |
430 | 496 | class TestJavaSDKInterfaceExample: |
431 | 497 | """Verify tasks written against ``InputTask`` receive their bound arguments.""" |
432 | 498 |
|
|
0 commit comments