cleveragents-core/features/steps/decision_dependency_basic_steps.py

"""Basic steps for decision dependency persistence feature."""

from __future__ import annotations

from behave import given, then, when
from behave.runner import Context

from cleveragents.core.exceptions import DatabaseError
from cleveragents.infrastructure.database.models import DecisionDependencyModel

from features.steps.helpers.decision_dependency_helpers import (
    PLAN_ID_1,
    PLAN_ID_2,
    create_child,
    create_prerequisite_action,
    create_prerequisite_plan,
    create_root,
    get_edges,
    record_dependency,
    setup_db,
)


@given("a fresh database for decision dependencies")
def step_fresh_db(context: Context) -> None:
    """Set up a fresh database."""
    setup_db(context)


@given('a prerequisite action "{action_name}" exists for dependencies')
def step_prereq_action(context: Context, action_name: str) -> None:
    """Create a prerequisite action."""
    create_prerequisite_action(context, action_name)


@given("a prerequisite plan exists for dependencies")
def step_prereq_plan(context: Context) -> None:
    """Create a prerequisite plan."""
    create_prerequisite_plan(context, PLAN_ID_1)


@given("a persisted root decision for dependencies")
def step_persisted_root_dep(context: Context) -> None:
    """Create and persist a root decision."""
    root = create_root(context)
    context._ddp_root_id = root.decision_id
    context._ddp_root_decision = root


@given("a persisted child decision for dependencies")
def step_persisted_child_dep(context: Context) -> None:
    """Create and persist a child decision."""
    root_id = getattr(context, "_ddp_root_id", None)
    if root_id is None:
        root_id = context._ddp_root_ids[0]
    child = create_child(context, root_id)
    context._ddp_child_id = child.decision_id
    context._ddp_child_decision = child


@given("{count:d} persisted sibling decisions for dependencies")
def step_persisted_siblings_dep(context: Context, count: int) -> None:
    """Create and persist multiple sibling decisions."""
    context._ddp_sibling_ids = []
    for index in range(count):
        sibling = create_child(
            context,
            context._ddp_root_id,
            sequence_number=index + 1,
        )
        context._ddp_sibling_ids.append(sibling.decision_id)


@given("{count:d} persisted root decisions for dependencies")
def step_persisted_roots_dep(context: Context, count: int) -> None:
    """Create and persist multiple root decisions."""
    context._ddp_root_ids = []
    for index in range(count):
        root = create_root(context, index=index)
        context._ddp_root_ids.append(root.decision_id)
    context._ddp_root_id = context._ddp_root_ids[0]


@given("{count:d} persisted plans for dependencies")
def step_persisted_plans_dep(context: Context, count: int) -> None:
    """Create and persist multiple plans."""
    context._ddp_plan_ids = [PLAN_ID_1]
    if count > 1:
        create_prerequisite_plan(context, PLAN_ID_2)
        context._ddp_plan_ids.append(PLAN_ID_2)


@given("a decision tree in plan 1 with dependencies")
def step_decision_tree_plan1_dep(context: Context) -> None:
    """Create a decision tree in plan 1 with dependencies."""
    root = create_root(context, PLAN_ID_1)
    child = create_child(context, root.decision_id, PLAN_ID_1)
    record_dependency(context, root.decision_id, child.decision_id)
    context._ddp_plan1_root_id = root.decision_id
    context._ddp_plan1_child_id = child.decision_id


@given("a decision tree in plan 2 with dependencies")
def step_decision_tree_plan2_dep(context: Context) -> None:
    """Create a decision tree in plan 2 with dependencies."""
    root = create_root(context, PLAN_ID_2)
    child = create_child(context, root.decision_id, PLAN_ID_2)
    record_dependency(context, root.decision_id, child.decision_id)
    context._ddp_plan2_root_id = root.decision_id
    context._ddp_plan2_child_id = child.decision_id


@given("a decision in plan 1 and a decision in plan 2")
def step_cross_plan_decisions_dep(context: Context) -> None:
    """Create decisions in different plans."""
    d1 = create_root(context, PLAN_ID_1)
    d2 = create_root(context, PLAN_ID_2)
    context._ddp_cross_plan_id1 = d1.decision_id
    context._ddp_cross_plan_id2 = d2.decision_id


@given("a recorded dependency from root to child")
def step_given_recorded_dependency_root_child(context: Context) -> None:
    """Record a dependency from root to child during scenario setup."""
    record_dependency(context, context._ddp_root_id, context._ddp_child_id)


@given("recorded dependencies from root to both siblings")
def step_given_recorded_dependencies_root_siblings(context: Context) -> None:
    """Record dependencies from root to all siblings during setup."""
    for sibling_id in context._ddp_sibling_ids:
        record_dependency(context, context._ddp_root_id, sibling_id)


@given("recorded dependencies from both roots to the child")
def step_given_recorded_dependencies_both_roots(context: Context) -> None:
    """Record dependencies from each root to the child during setup."""
    for root_id in context._ddp_root_ids:
        record_dependency(context, root_id, context._ddp_child_id)


@when("I record a dependency from the root to the child")
def step_record_dependency_root_child(context: Context) -> None:
    """Record a dependency from root to child."""
    record_dependency(context, context._ddp_root_id, context._ddp_child_id)


@when("I record dependencies from the root to both siblings")
def step_record_dependencies_root_siblings(context: Context) -> None:
    """Record dependencies from root to all siblings."""
    for sibling_id in context._ddp_sibling_ids:
        record_dependency(context, context._ddp_root_id, sibling_id)


@when('I record a dependency with relationship_type "{rel_type}"')
def step_record_dependency_custom_type(context: Context, rel_type: str) -> None:
    """Record a dependency with custom relationship type."""
    record_dependency(
        context,
        context._ddp_root_id,
        context._ddp_child_id,
        relationship_type=rel_type,
    )
    context._ddp_custom_rel_type = rel_type


@when("I try to record a self-loop dependency")
def step_try_self_loop_dependency(context: Context) -> None:
    """Try to record a self-loop dependency."""
    try:
        record_dependency(context, context._ddp_root_id, context._ddp_root_id)
        context._ddp_self_loop_error = None
    except DatabaseError as exc:
        context._ddp_self_loop_error = exc


@when("I try to record the same dependency again")
def step_try_duplicate_dependency(context: Context) -> None:
    """Try to record the same dependency again."""
    try:
        record_dependency(context, context._ddp_root_id, context._ddp_child_id)
        context._ddp_duplicate_error = None
    except DatabaseError as exc:
        context._ddp_duplicate_error = exc


@when("I try to record a cross-plan dependency")
def step_try_cross_plan_dependency(context: Context) -> None:
    """Try to record a cross-plan dependency."""
    try:
        record_dependency(
            context,
            context._ddp_cross_plan_id1,
            context._ddp_cross_plan_id2,
        )
        context._ddp_cross_plan_error = None
    except DatabaseError as exc:
        context._ddp_cross_plan_error = exc


@when("I get influence edges for the plan")
def step_get_influence_edges(context: Context) -> None:
    """Get influence edges for the plan."""
    context._ddp_influence_edges = get_edges(context, PLAN_ID_1)


@when("I get influence edges for plan 1")
def step_get_influence_edges_plan1(context: Context) -> None:
    """Get influence edges for plan 1."""
    context._ddp_influence_edges_plan1 = get_edges(context, PLAN_ID_1)


@when("I get influence edges for plan 2")
def step_get_influence_edges_plan2(context: Context) -> None:
    """Get influence edges for plan 2."""
    context._ddp_influence_edges_plan2 = get_edges(context, PLAN_ID_2)


@then("the dependency should be persisted to the database")
def step_assert_dependency_persisted(context: Context) -> None:
    """Assert that the dependency was persisted."""
    edges = get_edges(context)
    target_id = getattr(context, "_ddp_new_decision_id", None)
    if target_id is None:
        target_id = context._ddp_child_id
    assert context._ddp_root_id in edges, "Root should be in edges"
    assert target_id in edges[context._ddp_root_id], "Child should be in root's targets"


@then('the dependency should have relationship_type "{rel_type}"')
def step_assert_dependency_rel_type(context: Context, rel_type: str) -> None:
    """Assert that the dependency has the correct relationship type."""
    _assert_dependency_relationship_type(context, rel_type)


@then('the dependency should be persisted with relationship_type "{rel_type}"')
def step_assert_dependency_persisted_with_type(context: Context, rel_type: str) -> None:
    """Assert that the dependency was persisted with the correct type."""
    _assert_dependency_relationship_type(context, rel_type)


@then("both dependencies should be persisted to the database")
def step_assert_both_dependencies_persisted(context: Context) -> None:
    """Assert that both dependencies were persisted."""
    edges = get_edges(context)
    if hasattr(context, "_ddp_new_decision_id"):
        for root_id in context._ddp_root_ids:
            assert root_id in edges, f"Root {root_id} should be in edges"
            assert context._ddp_new_decision_id in edges[root_id], (
                f"New decision should be in {root_id}'s targets"
            )
        return

    assert context._ddp_root_id in edges, "Root should be in edges"
    assert len(edges[context._ddp_root_id]) == 2, "Root should have 2 targets"
    for sibling_id in context._ddp_sibling_ids:
        assert sibling_id in edges[context._ddp_root_id], (
            f"Sibling {sibling_id} should be in root's targets"
        )


@then("the influence edges should be empty")
def step_assert_influence_edges_empty(context: Context) -> None:
    """Assert that influence edges are empty."""
    assert context._ddp_influence_edges == {}, "Influence edges should be empty"


@then("the influence edges should contain {count:d} source")
def step_assert_influence_edges_source_count(context: Context, count: int) -> None:
    """Assert the number of sources in influence edges."""
    assert len(context._ddp_influence_edges) == count, (
        f"Should have {count} source(s), got {len(context._ddp_influence_edges)}"
    )


@then("the influence edges should contain {count:d} sources")
def step_assert_influence_edges_sources_count(context: Context, count: int) -> None:
    """Assert the number of sources in influence edges."""
    step_assert_influence_edges_source_count(context, count)


@then("the influence edges should map root to {targets}")
def step_assert_influence_edges_mapping(context: Context, targets: str) -> None:
    """Assert that influence edges map root to the specified targets."""
    if targets == "[child]":
        expected = [context._ddp_child_id]
    elif targets == "[sibling1, sibling2]":
        expected = context._ddp_sibling_ids
    else:
        expected = []

    assert context._ddp_root_id in context._ddp_influence_edges, (
        "Root should be in edges"
    )
    actual = context._ddp_influence_edges[context._ddp_root_id]
    assert set(actual) == set(expected), f"Expected {expected}, got {actual}"


@then("the influence edges should only include decisions from plan 1")
def step_assert_influence_edges_plan1_only(context: Context) -> None:
    """Assert that influence edges only include plan 1 decisions."""
    edges = context._ddp_influence_edges_plan1
    plan1_ids = {
        context._ddp_plan1_root_id,
        context._ddp_plan1_child_id,
    }
    for source, targets in edges.items():
        assert source in plan1_ids, f"Source {source} should be in plan 1"
        for target in targets:
            assert target in plan1_ids, f"Target {target} should be in plan 1"


@then("the cross-plan dependency should not appear in influence edges")
def step_assert_no_cross_plan_dependency(context: Context) -> None:
    """Assert that cross-plan dependencies don't appear in edges."""
    edges = get_edges(context)
    if context._ddp_cross_plan_id1 in edges:
        assert context._ddp_cross_plan_id2 not in edges[context._ddp_cross_plan_id1], (
            "Cross-plan dependency should not appear"
        )


@then("a database constraint error should be raised")
def step_assert_constraint_error(context: Context) -> None:
    """Assert that a constraint error was raised."""
    assert context._ddp_self_loop_error is not None, (
        "Should have raised a constraint error"
    )


@then("the operation should complete without error")
def step_assert_no_error(context: Context) -> None:
    """Assert that the operation completed without error."""
    assert context._ddp_duplicate_error is None, "Should not have raised an error"


def _assert_dependency_relationship_type(context: Context, rel_type: str) -> None:
    session = context._ddp_factory()
    dep = (
        session.query(DecisionDependencyModel)
        .filter_by(
            source_decision_id=context._ddp_root_id,
            target_decision_id=context._ddp_child_id,
        )
        .first()
    )
    assert dep is not None, "Dependency should exist"
    assert dep.relationship_type == rel_type, (
        f"Relationship type should be {rel_type}, got {dep.relationship_type}"
    )