pydantic--pydantic-ai/tests/graph/builder/test_graph_edge_cases.py

"""Additional edge case tests for graph execution to improve coverage."""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Literal

import pytest

from pydantic_graph import GraphBuilder, StepContext
from pydantic_graph.join import ReduceFirstValue, ReducerContext, reduce_sum

from ..._inline_snapshot import snapshot

pytestmark = pytest.mark.anyio


@dataclass
class MyState:
    value: int = 0


async def test_graph_repr():
    """Test that Graph.__repr__ returns a mermaid diagram."""
    g = GraphBuilder(state_type=MyState, output_type=int)

    @g.step
    async def simple_step(ctx: StepContext[MyState, None, None]) -> int:
        return 42  # pragma: no cover

    g.add(
        g.edge_from(g.start_node).to(simple_step),
        g.edge_from(simple_step).to(g.end_node),
    )

    graph = g.build()
    graph_repr = repr(graph)

    # Replace the non-constant graph object id with a constant string:
    normalized_graph_repr = re.sub(hex(id(graph)), '0xGraphObjectId', graph_repr)

    assert normalized_graph_repr == snapshot("""\
<pydantic_graph.graph_builder.Graph object at 0xGraphObjectId
stateDiagram-v2
  simple_step

  [*] --> simple_step
  simple_step --> [*]
>\
""")


async def test_graph_render_with_title():
    """Test Graph.render method with title parameter."""
    g = GraphBuilder(state_type=MyState, output_type=int)

    @g.step
    async def simple_step(ctx: StepContext[MyState, None, None]) -> int:
        return 42  # pragma: no cover

    g.add(
        g.edge_from(g.start_node).to(simple_step),
        g.edge_from(simple_step).to(g.end_node),
    )

    graph = g.build()
    rendered = graph.render(title='My Graph')
    assert rendered == snapshot("""\
---
title: My Graph
---
stateDiagram-v2
  simple_step

  [*] --> simple_step
  simple_step --> [*]\
""")


async def test_get_parent_fork_missing():
    """Test that get_parent_fork raises RuntimeError when join has no parent fork."""
    from pydantic_graph.id_types import JoinID, NodeID

    g = GraphBuilder(state_type=MyState, output_type=int)

    @g.step
    async def simple_step(ctx: StepContext[MyState, None, None]) -> int:
        return 42  # pragma: no cover

    g.add(
        g.edge_from(g.start_node).to(simple_step),
        g.edge_from(simple_step).to(g.end_node),
    )

    graph = g.build()

    # Try to get a parent fork for a non-existent join
    fake_join_id = JoinID(NodeID('fake_join'))
    with pytest.raises(RuntimeError, match='not a join node'):
        graph.get_parent_fork(fake_join_id)


async def test_decision_no_matching_branch():
    """Test that decision raises RuntimeError when no branch matches."""
    g = GraphBuilder(state_type=MyState, output_type=str)

    @g.step
    async def return_unexpected(ctx: StepContext[MyState, None, None]) -> int:
        return 999

    @g.step
    async def handle_str(ctx: StepContext[MyState, None, str]) -> str:
        return f'Got: {ctx.inputs}'  # pragma: no cover

    # the purpose of this test is to test runtime behavior when you have this type failure, which is why
    # we have the `# type: ignore` below
    g.add(
        g.edge_from(g.start_node).to(return_unexpected),
        g.edge_from(return_unexpected).to(g.decision().branch(g.match(str).to(handle_str))),  # type: ignore
        g.edge_from(handle_str).to(g.end_node),
    )

    graph = g.build()

    with pytest.raises(RuntimeError, match='No branch matched'):
        await graph.run(state=MyState())


async def test_decision_invalid_type_check():
    """Test decision branch with invalid type for isinstance check."""

    g = GraphBuilder(state_type=MyState, output_type=str)

    @g.step
    async def return_value(ctx: StepContext[MyState, None, None]) -> int:
        return 42

    @g.step
    async def handle_value(ctx: StepContext[MyState, None, int]) -> str:
        return str(ctx.inputs)

    # Try to use a non-type as a branch source - this might cause TypeError during isinstance check
    # Note: This is hard to trigger without directly constructing invalid decision branches
    # For now, just test normal union types work
    g.add(
        g.edge_from(g.start_node).to(return_value),
        g.edge_from(return_value).to(g.decision().branch(g.match(int).to(handle_value))),
        g.edge_from(handle_value).to(g.end_node),
    )

    graph = g.build()
    result = await graph.run(state=MyState())
    assert result == '42'


async def test_map_non_iterable():
    """Test that mapping a non-iterable value raises RuntimeError."""
    g = GraphBuilder(state_type=MyState, output_type=int)

    @g.step
    async def return_non_iterable(ctx: StepContext[MyState, None, None]) -> int:
        return 42  # Not iterable!

    @g.step
    async def process_item(ctx: StepContext[MyState, None, int]) -> int:
        return ctx.inputs  # pragma: no cover

    sum_items = g.join(reduce_sum, initial=0)

    # This will fail at runtime because we're trying to map over a non-iterable
    # We have a `# type: ignore` below because we are testing behavior when you ignore the type error
    g.add(
        g.edge_from(g.start_node).to(return_non_iterable),
        g.edge_from(return_non_iterable).map().to(process_item),  # type: ignore
        g.edge_from(process_item).to(sum_items),
        g.edge_from(sum_items).to(g.end_node),
    )

    graph = g.build()

    with pytest.raises(RuntimeError, match='Cannot map non-iterable'):
        await graph.run(state=MyState())


async def test_reducer_stop_iteration():
    """Test reducer that raises StopIteration to cancel concurrent tasks."""

    @dataclass
    class EarlyStopState:
        stopped: bool = False

    g = GraphBuilder(state_type=EarlyStopState, output_type=int)

    @g.step
    async def generate_numbers(ctx: StepContext[EarlyStopState, None, None]) -> list[int]:
        return [1, 2, 3, 4, 5]

    @g.step
    async def slow_process(ctx: StepContext[EarlyStopState, None, int]) -> int:
        # Simulate some processing
        return ctx.inputs * 2

    def get_early_stopping_reducer():
        count = 0

        def reduce(ctx: ReducerContext[EarlyStopState, object], current: int, inputs: int) -> int:
            nonlocal count
            count += 1
            current += inputs
            if count >= 2:
                ctx.state.stopped = True  # update the state so we can assert on it later
                ctx.cancel_sibling_tasks()
            return current

        return reduce

    stop_early = g.join(get_early_stopping_reducer(), initial=0)

    g.add(
        g.edge_from(g.start_node).to(generate_numbers),
        g.edge_from(generate_numbers).map().to(slow_process),
        g.edge_from(slow_process).to(stop_early),
        g.edge_from(stop_early).to(g.end_node),
    )

    graph = g.build()
    state = EarlyStopState()
    result = await graph.run(state=state)

    # Should have stopped early
    assert state.stopped
    # Result should be less than the full sum (2+4+6+8+10=30)
    # Actually, it should be less than the maximum of any two terms, (8+10=18)
    assert result <= 18


async def test_parallel_reducer_stop_iteration_explicit_fork_ids():
    """Test reducer that raises StopIteration to cancel concurrent tasks."""

    g = GraphBuilder(output_type=int)

    @g.step
    async def generate_numbers(ctx: StepContext[None, None, None]) -> list[int]:
        return [1, 1, 1, 1, 1]

    stop_early_1 = g.join(ReduceFirstValue[int](), initial=0, parent_fork_id='map_1')
    stop_early_2 = g.join(ReduceFirstValue[int](), initial=0, parent_fork_id='map_2')
    collect = g.join(reduce_sum, initial=0)

    g.add(
        g.edge_from(g.start_node).to(generate_numbers),
        g.edge_from(generate_numbers).broadcast(
            lambda b: [
                b.map(fork_id='map_1').transform(lambda ctx: ctx.inputs * 10).to(stop_early_1),
                b.map(fork_id='map_2').to(stop_early_2),
            ]
        ),
        g.edge_from(stop_early_1, stop_early_2).to(collect),
        g.edge_from(collect).to(g.end_node),
    )

    graph = g.build()
    result = await graph.run()

    # Result should be 11 because it should have stopped early on one input in the first fork and the *10 fork
    assert result == 11


async def test_parallel_reducer_stop_iteration_implicit_fork_ids():
    """Test reducer that raises StopIteration to cancel concurrent tasks."""

    g = GraphBuilder(output_type=int)

    @g.step
    async def generate_numbers(ctx: StepContext[None, None, None]) -> list[int]:
        return [1, 1, 1, 1, 1]

    stop_early_1 = g.join(ReduceFirstValue[int](), initial=0, preferred_parent_fork='closest')
    stop_early_2 = g.join(ReduceFirstValue[int](), initial=0, preferred_parent_fork='closest')
    collect = g.join(reduce_sum, initial=0)

    g.add(
        g.edge_from(g.start_node).to(generate_numbers),
        g.edge_from(generate_numbers).broadcast(
            lambda b: [
                b.map().transform(lambda ctx: ctx.inputs * 10).to(stop_early_1),
                b.map().to(stop_early_2),
            ]
        ),
        g.edge_from(stop_early_1, stop_early_2).to(collect),
        g.edge_from(collect).to(g.end_node),
    )

    graph = g.build()
    result = await graph.run()

    # Result should be 11 because it should have stopped early on one input in the first fork and the *10 fork
    assert result == 11


async def test_empty_path_handling():
    """Test handling of empty paths in graph execution."""
    g = GraphBuilder(state_type=MyState, output_type=int)

    @g.step
    async def return_value(ctx: StepContext[MyState, None, None]) -> int:
        return 42

    # Just connect start to step to end - this should work fine
    g.add(
        g.edge_from(g.start_node).to(return_value),
        g.edge_from(return_value).to(g.end_node),
    )

    graph = g.build()
    result = await graph.run(state=MyState())
    assert result == 42


async def test_literal_branch_matching():
    """Test decision branch matching with Literal types."""
    g = GraphBuilder(state_type=MyState, output_type=str)

    @g.step
    async def choose_option(ctx: StepContext[MyState, None, None]) -> Literal['a', 'b', 'c']:
        return 'b'

    @g.step
    async def handle_a(ctx: StepContext[MyState, None, object]) -> str:
        return 'Chose A'  # pragma: no cover

    @g.step
    async def handle_b(ctx: StepContext[MyState, None, object]) -> str:
        return 'Chose B'

    @g.step
    async def handle_c(ctx: StepContext[MyState, None, object]) -> str:
        return 'Chose C'  # pragma: no cover

    from pydantic_graph import TypeExpression

    g.add(
        g.edge_from(g.start_node).to(choose_option),
        g.edge_from(choose_option).to(
            g.decision()
            .branch(g.match(TypeExpression[Literal['a']]).to(handle_a))
            .branch(g.match(TypeExpression[Literal['b']]).to(handle_b))
            .branch(g.match(TypeExpression[Literal['c']]).to(handle_c))
        ),
        g.edge_from(handle_a, handle_b, handle_c).to(g.end_node),
    )

    graph = g.build()
    result = await graph.run(state=MyState())
    assert result == 'Chose B'


async def test_path_with_label_marker():
    """Test that LabelMarker in paths doesn't affect execution."""
    g = GraphBuilder(state_type=MyState, output_type=int)

    @g.step
    async def step_a(ctx: StepContext[MyState, None, None]) -> int:
        return 10

    @g.step
    async def step_b(ctx: StepContext[MyState, None, int]) -> int:
        return ctx.inputs * 2

    # Add labels to the path
    g.add(
        g.edge_from(g.start_node).label('start').to(step_a),
        g.edge_from(step_a).label('middle').to(step_b),
        g.edge_from(step_b).label('end').to(g.end_node),
    )

    graph = g.build()
    result = await graph.run(state=MyState())
    assert result == 20


async def test_nested_reducers_with_prefix():
    """Test multiple active reducers where one is a prefix of another."""
    g = GraphBuilder(state_type=MyState, output_type=int)

    @g.step
    async def outer_list(ctx: StepContext[MyState, None, None]) -> list[list[int]]:
        return [[1, 2], [3, 4]]

    @g.step
    async def inner_process(ctx: StepContext[MyState, None, int]) -> int:
        return ctx.inputs * 2

    # Note: we use  the _most_ ancestral fork as the parent fork by default, which means that this join
    # actually will join all forks from the initial outer_list, therefore summing everything, rather
    # than _only_ summing the inner loops. If/when we add more control over the parent fork calculation, we can
    # test that it's possible to use separate logic for the inside vs. the outside.
    sum_join = g.join(reduce_sum, initial=0)

    # Create nested map operations
    g.add(
        g.edge_from(g.start_node).to(outer_list),
        g.edge_from(outer_list).map().map().to(inner_process),
        g.edge_from(inner_process).to(sum_join),
        g.edge_from(sum_join).to(g.end_node),
    )

    graph = g.build()
    result = await graph.run(state=MyState())
    # (1+2+3+4) * 2 = 20
    assert result == 20
    assert str(graph) == snapshot("""\
stateDiagram-v2
  outer_list
  state map <<fork>>
  state map_2 <<fork>>
  inner_process
  state reduce_sum <<join>>

  [*] --> outer_list
  outer_list --> map
  map --> map_2
  map_2 --> inner_process
  inner_process --> reduce_sum
  reduce_sum --> [*]\
""")