cleveragents-core/features/steps/provider_shared_steps.py

from __future__ import annotations

import ast
from typing import Any
from unittest.mock import MagicMock, patch

from behave import given

from cleveragents.domain.models.core import (
    Context,
    OperationType,
    Plan,
    Project,
)


def _register_cleanup(context, cleanup):
    if hasattr(context, "add_cleanup"):
        context.add_cleanup(cleanup)
    else:
        cleanup_handlers = getattr(context, "_cleanup_handlers", [])
        cleanup_handlers.append(cleanup)
        context._cleanup_handlers = cleanup_handlers


def _parse_kwargs_string(kwargs_string: str) -> dict[str, Any]:
    if not kwargs_string:
        return {}
    result: dict[str, Any] = {}
    for entry in kwargs_string.split(","):
        entry = entry.strip()
        if not entry:
            continue
        if "=" not in entry:
            result[entry] = True
            continue
        key, value = entry.split("=", 1)
        key = key.strip()
        value = value.strip()
        try:
            parsed_value = ast.literal_eval(value)
        except Exception:
            parsed_value = value
        result[key] = parsed_value
    return result


def _setup_plan_generation_graph(context) -> MagicMock:
    patcher = patch(
        "cleveragents.providers.llm.langchain_chat_provider.PlanGenerationGraph"
    )
    context.plan_generation_patcher = patcher
    mock_graph_class = patcher.start()
    _register_cleanup(context, patcher.stop)
    mock_graph_instance = MagicMock(name="PlanGenerationGraphInstance")

    default_state = {
        "generated_changes": [],
        "validation_result": {"status": "PASS"},
        "error": None,
    }
    state_override = getattr(context, "plan_generation_state_override", None)
    mock_graph_instance.invoke.return_value = state_override or default_state

    invoke_side_effect = getattr(context, "plan_generation_invoke_side_effect", None)
    if invoke_side_effect is not None:
        mock_graph_instance.invoke.side_effect = invoke_side_effect

    stream_events = getattr(context, "plan_generation_stream_events", None)
    if stream_events is None:
        mock_graph_instance.stream.return_value = iter(())
    else:
        events_copy = list(stream_events)

        def _stream(*_args, **_kwargs):
            yield from events_copy

        mock_graph_instance.stream.side_effect = _stream

    mock_graph_class.return_value = mock_graph_instance
    context.plan_generation_graph = mock_graph_instance
    context.plan_generation_graph_class = mock_graph_class
    return mock_graph_instance


@given("I have sample provider domain inputs")
def step_sample_provider_inputs(context):
    context.project = MagicMock(spec=Project)
    context.plan = MagicMock(spec=Plan)
    context.plan.prompt = "Add placeholder coverage"
    context.contexts = [MagicMock(spec=Context)]
    context.contexts[0].content = "Sample context entry"


@given('the plan generation graph returns a generated change for "{file_path}"')
def step_plan_generation_returns_change(context, file_path):
    change = {
        "plan_id": 1,
        "file_path": file_path,
        "operation": OperationType.MODIFY.value,
        "new_content": "# updated content",
    }
    context.plan_generation_state_override = {
        "generated_changes": [change],
        "validation_result": {"status": "PASS"},
        "error": None,
    }


@given('the plan generation graph emits streaming nodes "{node_list}"')
def step_plan_generation_stream_nodes(context, node_list):
    nodes = [node.strip() for node in node_list.split(",") if node.strip()]
    state = getattr(context, "plan_generation_state_override", {}) or {}
    events: list[dict[str, Any]] = []
    for node in nodes:
        payload: dict[str, Any] = {"status": "completed"}
        if node == "generate_plan" and isinstance(state.get("generated_changes"), list):
            payload = {
                "generated_changes": state["generated_changes"],
                "status": "completed",
            }
        elif node == "validate" and isinstance(state.get("validation_result"), dict):
            payload = {
                "validation_result": state["validation_result"],
                "status": "completed",
            }
        events.append({node: payload})
    context.plan_generation_stream_events = events


@given('the plan generation graph raises ValueError "{message}"')
def step_plan_generation_raises_value_error(context, message):
    context.plan_generation_invoke_side_effect = ValueError(message)


@given('the plan generation graph raises RuntimeError "{message}"')
def step_plan_generation_runtime_error(context, message):
    context.plan_generation_invoke_side_effect = RuntimeError(message)