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)