|
14 | 14 |
|
15 | 15 | """Testings for the JoinNode.""" |
16 | 16 |
|
| 17 | +import asyncio |
| 18 | + |
17 | 19 | from google.adk import workflow |
18 | 20 | from google.adk.apps import app |
| 21 | +from google.adk.apps.app import ResumabilityConfig |
| 22 | +from google.adk.events._interrupts import extract_event_interrupt_ids |
| 23 | +from google.adk.events.event import Event |
| 24 | +from google.adk.events.request_input import RequestInput |
19 | 25 | from google.adk.workflow import _base_node as base_node |
20 | 26 | from google.adk.workflow import _graph as workflow_graph |
21 | 27 | from google.adk.workflow import _join_node as join_node |
22 | 28 | from google.adk.workflow import START |
23 | 29 | from google.adk.workflow._workflow import Workflow |
| 30 | +from google.adk.workflow.utils import _workflow_hitl_utils as hitl_utils |
24 | 31 | from pydantic import BaseModel |
25 | 32 | import pytest |
26 | 33 |
|
@@ -348,3 +355,213 @@ async def test_join_node_computes_common_branch_prefix( |
348 | 355 | assert len(capture_events) > 0 |
349 | 356 | for e in capture_events: |
350 | 357 | assert e.branch is None |
| 358 | + |
| 359 | + |
| 360 | +@pytest.mark.asyncio |
| 361 | +async def test_join_node_in_loop_waits_for_each_iteration( |
| 362 | + request: pytest.FixtureRequest, |
| 363 | +) -> None: |
| 364 | + """Tests JoinNode in a loop waits for every predecessor's new run.""" |
| 365 | + join_inputs: list[dict[str, str]] = [] |
| 366 | + |
| 367 | + def node_a() -> str: |
| 368 | + return 'A2' if join_inputs else 'A1' |
| 369 | + |
| 370 | + def node_b() -> str: |
| 371 | + return 'B2' if join_inputs else 'B1' |
| 372 | + |
| 373 | + def node_check(node_input: dict[str, str]) -> Event | None: |
| 374 | + join_inputs.append(node_input) |
| 375 | + return Event(route='again') if len(join_inputs) == 1 else None |
| 376 | + |
| 377 | + join = join_node.JoinNode(name='join') |
| 378 | + agent = Workflow( |
| 379 | + name='wf', |
| 380 | + edges=[ |
| 381 | + (START, node_a, node_b), |
| 382 | + ((node_a, node_b), join, node_check), |
| 383 | + (node_check, {'again': node_a}), |
| 384 | + ], |
| 385 | + ) |
| 386 | + runner = testing_utils.InMemoryRunner( |
| 387 | + app=app.App(name=request.function.__name__, root_agent=agent) |
| 388 | + ) |
| 389 | + |
| 390 | + await runner.run_async('start') |
| 391 | + |
| 392 | + assert join_inputs == [ |
| 393 | + {'node_a': 'A1', 'node_b': 'B1'}, |
| 394 | + {'node_a': 'A2', 'node_b': 'B2'}, |
| 395 | + ] |
| 396 | + |
| 397 | + |
| 398 | +@pytest.mark.asyncio |
| 399 | +@pytest.mark.parametrize('is_resumable', [False, True]) |
| 400 | +async def test_join_node_in_loop_waits_for_predecessor_paused_for_input( |
| 401 | + request: pytest.FixtureRequest, is_resumable: bool |
| 402 | +) -> None: |
| 403 | + """Tests JoinNode in a loop waits for a predecessor that asks the user.""" |
| 404 | + join_inputs: list[dict[str, str]] = [] |
| 405 | + |
| 406 | + def node_a() -> str: |
| 407 | + return 'A2' if join_inputs else 'A1' |
| 408 | + |
| 409 | + def node_b() -> RequestInput | str: |
| 410 | + return RequestInput(message='Continue?') if join_inputs else 'B1' |
| 411 | + |
| 412 | + def node_check(node_input: dict[str, str]) -> Event | None: |
| 413 | + join_inputs.append(node_input) |
| 414 | + return Event(route='again') if len(join_inputs) == 1 else None |
| 415 | + |
| 416 | + join = join_node.JoinNode(name='join') |
| 417 | + agent = Workflow( |
| 418 | + name='wf', |
| 419 | + edges=[ |
| 420 | + (START, node_a, node_b), |
| 421 | + ((node_a, node_b), join, node_check), |
| 422 | + (node_check, {'again': node_a}), |
| 423 | + ], |
| 424 | + ) |
| 425 | + runner = testing_utils.InMemoryRunner( |
| 426 | + app=app.App( |
| 427 | + name=request.function.__name__, |
| 428 | + root_agent=agent, |
| 429 | + resumability_config=ResumabilityConfig(is_resumable=is_resumable), |
| 430 | + ) |
| 431 | + ) |
| 432 | + |
| 433 | + events = await runner.run_async('start') |
| 434 | + |
| 435 | + assert join_inputs == [{'node_a': 'A1', 'node_b': 'B1'}] |
| 436 | + |
| 437 | + request_events = workflow_testing_utils.get_request_input_events(events) |
| 438 | + [interrupt_id] = extract_event_interrupt_ids(request_events[0]) |
| 439 | + await runner.run_async( |
| 440 | + new_message=hitl_utils.create_request_input_response( |
| 441 | + interrupt_id, {'result': 'B2'} |
| 442 | + ), |
| 443 | + invocation_id=events[0].invocation_id, |
| 444 | + ) |
| 445 | + |
| 446 | + assert join_inputs == [ |
| 447 | + {'node_a': 'A1', 'node_b': 'B1'}, |
| 448 | + {'node_a': 'A2', 'node_b': 'B2'}, |
| 449 | + ] |
| 450 | + |
| 451 | + |
| 452 | +@pytest.mark.asyncio |
| 453 | +async def test_join_node_in_loop_clears_prior_output_when_node_returns_none( |
| 454 | + request: pytest.FixtureRequest, |
| 455 | +) -> None: |
| 456 | + """Tests JoinNode receives None when a rerun predecessor returns None.""" |
| 457 | + join_inputs: list[dict[str, str | None]] = [] |
| 458 | + |
| 459 | + def node_a() -> str | None: |
| 460 | + return None if join_inputs else 'A1' |
| 461 | + |
| 462 | + def node_b() -> str: |
| 463 | + return 'B2' if join_inputs else 'B1' |
| 464 | + |
| 465 | + def node_check(node_input: dict[str, str | None]) -> Event | None: |
| 466 | + join_inputs.append(node_input) |
| 467 | + return Event(route='again') if len(join_inputs) == 1 else None |
| 468 | + |
| 469 | + join = join_node.JoinNode(name='join') |
| 470 | + agent = Workflow( |
| 471 | + name='wf', |
| 472 | + edges=[ |
| 473 | + (START, node_a, node_b), |
| 474 | + ((node_a, node_b), join, node_check), |
| 475 | + (node_check, {'again': node_a}), |
| 476 | + ], |
| 477 | + ) |
| 478 | + runner = testing_utils.InMemoryRunner( |
| 479 | + app=app.App(name=request.function.__name__, root_agent=agent) |
| 480 | + ) |
| 481 | + |
| 482 | + await runner.run_async('start') |
| 483 | + |
| 484 | + assert join_inputs == [ |
| 485 | + {'node_a': 'A1', 'node_b': 'B1'}, |
| 486 | + {'node_a': None, 'node_b': 'B2'}, |
| 487 | + ] |
| 488 | + |
| 489 | + |
| 490 | +@pytest.mark.asyncio |
| 491 | +async def test_join_node_does_not_fire_when_predecessor_routes_away( |
| 492 | + request: pytest.FixtureRequest, |
| 493 | +) -> None: |
| 494 | + """Tests JoinNode does not fire when a predecessor routes away from it.""" |
| 495 | + |
| 496 | + def node_a() -> Event: |
| 497 | + return Event(output='A1', route='away') |
| 498 | + |
| 499 | + def node_b() -> str: |
| 500 | + return 'B1' |
| 501 | + |
| 502 | + join = join_node.JoinNode(name='join') |
| 503 | + other_node = workflow_testing_utils.InputCapturingNode(name='other_node') |
| 504 | + capture = workflow_testing_utils.InputCapturingNode(name='capture') |
| 505 | + agent = Workflow( |
| 506 | + name='wf', |
| 507 | + edges=[ |
| 508 | + (START, (node_a, node_b)), |
| 509 | + (node_a, {'to_join': join, 'away': other_node}), |
| 510 | + (node_b, join), |
| 511 | + (join, capture), |
| 512 | + ], |
| 513 | + ) |
| 514 | + runner = testing_utils.InMemoryRunner( |
| 515 | + app=app.App(name=request.function.__name__, root_agent=agent) |
| 516 | + ) |
| 517 | + |
| 518 | + await runner.run_async('start') |
| 519 | + |
| 520 | + assert other_node.received_inputs == ['A1'] |
| 521 | + assert capture.received_inputs == [] |
| 522 | + |
| 523 | + |
| 524 | +@pytest.mark.asyncio |
| 525 | +async def test_join_node_waits_for_retriggered_in_flight_predecessor( |
| 526 | + request: pytest.FixtureRequest, |
| 527 | +) -> None: |
| 528 | + """Tests JoinNode waits when an arrived predecessor has a run in flight.""" |
| 529 | + a_runs = 0 |
| 530 | + b_runs = 0 |
| 531 | + |
| 532 | + async def node_a() -> str: |
| 533 | + nonlocal a_runs |
| 534 | + a_runs += 1 |
| 535 | + if a_runs == 2: |
| 536 | + for _ in range(5): |
| 537 | + await asyncio.sleep(0) |
| 538 | + return f'A{a_runs}' |
| 539 | + |
| 540 | + def node_b() -> Event: |
| 541 | + nonlocal b_runs |
| 542 | + b_runs += 1 |
| 543 | + route = 'retry' if b_runs == 1 else 'to_join' |
| 544 | + return Event(output=f'B{b_runs}', route=route) |
| 545 | + |
| 546 | + def restart() -> None: |
| 547 | + return None |
| 548 | + |
| 549 | + join = join_node.JoinNode(name='join') |
| 550 | + capture = workflow_testing_utils.InputCapturingNode(name='capture') |
| 551 | + agent = Workflow( |
| 552 | + name='wf', |
| 553 | + edges=[ |
| 554 | + (START, (node_a, node_b)), |
| 555 | + (node_a, join), |
| 556 | + (node_b, {'to_join': join, 'retry': restart}), |
| 557 | + (restart, (node_a, node_b)), |
| 558 | + (join, capture), |
| 559 | + ], |
| 560 | + ) |
| 561 | + runner = testing_utils.InMemoryRunner( |
| 562 | + app=app.App(name=request.function.__name__, root_agent=agent) |
| 563 | + ) |
| 564 | + |
| 565 | + await runner.run_async('start') |
| 566 | + |
| 567 | + assert capture.received_inputs == [{'node_a': 'A2', 'node_b': 'B2'}] |
0 commit comments