cleveragents-core/features/steps/unified_context_models_steps.py

"""Step definitions for features/unified_context_models.feature.

Tests isinstance compatibility and construction of the unified
ContextFragment model hierarchy.
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta

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

from cleveragents.domain.models.acms.crp import (
    AssembledContext as CRPAssembledContext,
)
from cleveragents.domain.models.acms.crp import (
    ContextBudget as CRPContextBudget,
)
from cleveragents.domain.models.acms.crp import (
    ContextFragment as CRPContextFragment,
)
from cleveragents.domain.models.acms.crp import (
    FragmentProvenance as CRPFragmentProvenance,
)
from cleveragents.domain.models.core.context_fragment import (
    ContextBudget,
    ContextFragment,
    ContextPayload,
    FragmentProvenance,
)

# Default provenance for test fragments.
_DEFAULT_PROVENANCE = FragmentProvenance(resource_uri="test://unified")


# ---------------------------------------------------------------------------
# Given — Module availability
# ---------------------------------------------------------------------------


@given("the unified model modules are available")
def step_modules_available(context: Context) -> None:
    pass


# ---------------------------------------------------------------------------
# Given — FragmentProvenance
# ---------------------------------------------------------------------------


@given('a core FragmentProvenance with resource_uri "{uri}"')
def step_core_provenance(context: Context, uri: str) -> None:
    context.provenance = FragmentProvenance(resource_uri=uri)


@given('a core FragmentProvenance for strategy "{strategy}" with resource_uri "{uri}"')
def step_core_provenance_with_strategy(
    context: Context, strategy: str, uri: str
) -> None:
    context.provenance = FragmentProvenance(resource_uri=uri, strategy=strategy)


@given('a CRP FragmentProvenance with resource_uri "{uri}"')
def step_crp_provenance(context: Context, uri: str) -> None:
    context.crp_provenance = CRPFragmentProvenance(resource_uri=uri)


# ---------------------------------------------------------------------------
# Given — ContextFragment
# ---------------------------------------------------------------------------


@given('a core ContextFragment with uko_node "{uko_node}" and content "{content}"')
def step_core_fragment(context: Context, uko_node: str, content: str) -> None:
    context.unified_fragment = ContextFragment(
        uko_node=uko_node,
        content=content,
        token_count=len(content),
        provenance=_DEFAULT_PROVENANCE,
    )


@given("two core ContextFragments with identical fields")
def step_two_identical_fragments(context: Context) -> None:
    fixed_time = datetime.now(UTC)
    frag_a = ContextFragment(
        fragment_id="FIXEDIDFOREQUALITYTEST12345",
        uko_node="project://eq-test",
        content="equal",
        token_count=5,
        provenance=_DEFAULT_PROVENANCE,
        created_at=fixed_time,
    )
    frag_b = ContextFragment(
        fragment_id="FIXEDIDFOREQUALITYTEST12345",
        uko_node="project://eq-test",
        content="equal",
        token_count=5,
        provenance=_DEFAULT_PROVENANCE,
        created_at=fixed_time,
    )
    context.frag_a = frag_a
    context.frag_b = frag_b


# ---------------------------------------------------------------------------
# Given — ContextBudget
# ---------------------------------------------------------------------------


@given("a core ContextBudget with max_tokens {max_tok:d} and reserved_tokens {res:d}")
def step_core_budget(context: Context, max_tok: int, res: int) -> None:
    context.unified_budget = ContextBudget(max_tokens=max_tok, reserved_tokens=res)


# ---------------------------------------------------------------------------
# Given — ContextPayload
# ---------------------------------------------------------------------------


@given('a core ContextPayload with plan_id "{plan_id}"')
def step_core_payload(context: Context, plan_id: str) -> None:
    context.unified_payload = ContextPayload(plan_id=plan_id)


# ---------------------------------------------------------------------------
# When — Mutation attempts
# ---------------------------------------------------------------------------


@when("I attempt to mutate the fragment uko_node")
def step_mutate_fragment(context: Context) -> None:
    context.mutation_error = None
    try:
        context.unified_fragment.uko_node = "project://mutated"  # type: ignore[misc]
    except ValidationError as exc:
        context.mutation_error = exc


@when("I attempt to mutate the budget max_tokens")
def step_mutate_budget(context: Context) -> None:
    context.mutation_error = None
    try:
        context.unified_budget.max_tokens = 9999  # type: ignore[misc]
    except ValidationError as exc:
        context.mutation_error = exc


@when("I attempt to mutate the CRP provenance resource_uri")
def step_mutate_crp_provenance(context: Context) -> None:
    context.mutation_error = None
    try:
        context.crp_provenance.resource_uri = "test://mutated"  # type: ignore[misc]
    except ValidationError as exc:
        context.mutation_error = exc


@when("I create a core ContextBudget with reserved_tokens equal to max_tokens")
def step_create_budget_equal(context: Context) -> None:
    context.strict_budget_error = None
    try:
        ContextBudget(max_tokens=100, reserved_tokens=100)
    except ValidationError as exc:
        context.strict_budget_error = exc


# ---------------------------------------------------------------------------
# Then — isinstance checks
# ---------------------------------------------------------------------------


@then("the provenance should be an instance of CRP FragmentProvenance")
def step_provenance_isinstance(context: Context) -> None:
    assert isinstance(context.provenance, CRPFragmentProvenance), (
        f"Expected isinstance of CRP FragmentProvenance, got {type(context.provenance)}"
    )


@then("the fragment should be an instance of CRP ContextFragment")
def step_fragment_isinstance(context: Context) -> None:
    assert isinstance(context.unified_fragment, CRPContextFragment), (
        f"Expected isinstance of CRP ContextFragment, "
        f"got {type(context.unified_fragment)}"
    )


@then("the budget should be an instance of CRP ContextBudget")
def step_budget_isinstance(context: Context) -> None:
    assert isinstance(context.unified_budget, CRPContextBudget), (
        f"Expected isinstance of CRP ContextBudget, got {type(context.unified_budget)}"
    )


@then("the payload should be an instance of CRP AssembledContext")
def step_payload_isinstance(context: Context) -> None:
    assert isinstance(context.unified_payload, CRPAssembledContext), (
        f"Expected isinstance of CRP AssembledContext, "
        f"got {type(context.unified_payload)}"
    )


# ---------------------------------------------------------------------------
# Then — Field value assertions
# ---------------------------------------------------------------------------


@then('the provenance resource_type should be "{value}"')
def step_provenance_resource_type(context: Context, value: str) -> None:
    assert context.provenance.resource_type == value


@then('the provenance strategy should be "{value}"')
def step_provenance_strategy(context: Context, value: str) -> None:
    assert context.provenance.strategy == value


@then('the unified fragment tier should be "{tier}"')
def step_fragment_tier(context: Context, tier: str) -> None:
    assert context.unified_fragment.tier == tier


@then("the unified fragment fragment_id should be non-empty")
def step_fragment_id_nonempty(context: Context) -> None:
    assert context.unified_fragment.fragment_id
    assert len(context.unified_fragment.fragment_id) > 0


@then("the unified fragment fragment_id should be {length:d} characters")
def step_fragment_id_length(context: Context, length: int) -> None:
    fid = context.unified_fragment.fragment_id
    assert len(fid) == length, (
        f"Expected fragment_id of {length} chars, got {len(fid)}: {fid!r}"
    )


@then("the unified fragment created_at should be a recent datetime")
def step_fragment_created_at_recent(context: Context) -> None:
    created = context.unified_fragment.created_at
    assert isinstance(created, datetime)
    now = datetime.now(UTC)
    assert now - created < timedelta(seconds=60)


@then("the unified budget available_tokens should be {tokens:d}")
def step_budget_available(context: Context, tokens: int) -> None:
    assert context.unified_budget.available_tokens == tokens


# ---------------------------------------------------------------------------
# Then — Mutation error assertions
# ---------------------------------------------------------------------------


@then("a frozen mutation error should be raised")
def step_frozen_mutation_error(context: Context) -> None:
    assert context.mutation_error is not None, (
        "Expected a frozen mutation error but none was raised"
    )


# ---------------------------------------------------------------------------
# Then — Equality assertions
# ---------------------------------------------------------------------------


@then("the two fragments should be equal")
def step_fragments_equal(context: Context) -> None:
    assert context.frag_a == context.frag_b


@then("a validation error should be raised for strict budget")
def step_strict_budget_error(context: Context) -> None:
    assert context.strict_budget_error is not None, (
        "Expected a validation error for reserved == max but none was raised"
    )