import pytest from app.services import ai_tree_builder as atb class _FakeProvider: def __init__(self, raw): self._raw = raw async def generate_json(self, *, system_prompt, messages, max_tokens): return self._raw, None, None @pytest.mark.asyncio async def test_generate_next_node_assigns_id_when_model_omits_it(monkeypatch): """The SYSTEM_PROMPT never asks the model for an id (Finding 1). The server must assign one to every generated node, or the advance protocol — which keys on node_id — can never record an answer and the walk stalls on question 1.""" monkeypatch.setattr( atb, "get_ai_provider", lambda *a, **k: _FakeProvider('{"node_type":"question","text":"Plugged in?"}'), ) node = await atb.generate_next_node("printer down", "printer", []) assert node["node_type"] == "question" assert node.get("id"), "generated node must carry a server-assigned id" @pytest.mark.asyncio async def test_generate_next_node_depth_cap_node_has_id(monkeypatch): """The depth-cap escalate node must also carry an id (it is persisted as current_node_id and may be appended to walked_path).""" walked = [{"node_type": "question", "id": f"n{i}", "text": "?", "answer": "no"} for i in range(atb.MAX_DEPTH)] node = await atb.generate_next_node("x", "printer", walked) assert node["node_type"] == "escalate" assert node.get("id") @pytest.mark.asyncio async def test_generate_next_node_generation_failed_node_has_id(monkeypatch): """When both generation attempts fail, the fallback escalate node carries an id.""" monkeypatch.setattr( atb, "get_ai_provider", lambda *a, **k: _FakeProvider("not json at all"), ) node = await atb.generate_next_node("x", "printer", []) assert node["node_type"] == "escalate" assert node["reason_category"] == "generation_failed" assert node.get("id") # --------------------------------------------------------------------------- # Answer labels: the button text must match the question (live-walk defect: # "Microsoft account or local account?" rendered with Yes/No buttons). # --------------------------------------------------------------------------- def test_system_prompt_requires_answer_labels(): """The prompt must mandate yes_label/no_label on question nodes — the prompt forcing label-less '' output is the root cause of the question/button mismatch.""" assert "yes_label" in atb.SYSTEM_PROMPT and "no_label" in atb.SYSTEM_PROMPT @pytest.mark.asyncio async def test_generated_question_passes_labels_through(monkeypatch): monkeypatch.setattr( atb, "get_ai_provider", lambda *a, **k: _FakeProvider( '{"node_type":"question",' '"text":"Is Jane\'s Windows account a Microsoft account or a local account?",' '"yes_label":"Microsoft account","no_label":"Local account"}' ), ) node = await atb.generate_next_node("login issue", "account_login", []) assert node["yes_label"] == "Microsoft account" assert node["no_label"] == "Local account" @pytest.mark.asyncio async def test_question_missing_labels_gets_yes_no_defaults(monkeypatch): monkeypatch.setattr( atb, "get_ai_provider", lambda *a, **k: _FakeProvider('{"node_type":"question","text":"Is the printer powered on?"}'), ) node = await atb.generate_next_node("printer down", "printer", []) assert node["yes_label"] == "Yes" assert node["no_label"] == "No" def test_validate_node_rejects_hard_floor_text_in_labels(): node = {"node_type": "question", "text": "How should we proceed?", "yes_label": "Edit the registry", "no_label": "Wait"} with pytest.raises(atb.UnsafeNodeError): atb.validate_node(node) def test_validate_node_rejects_indistinct_or_malformed_labels(): base = {"node_type": "question", "text": "Which network is the laptop on?"} with pytest.raises(atb.UnsafeNodeError): atb.validate_node({**base, "yes_label": "Wi-Fi", "no_label": "wi-fi "}) with pytest.raises(atb.UnsafeNodeError): atb.validate_node({**base, "yes_label": 1, "no_label": "Ethernet"}) def test_build_context_prefers_answer_label_over_raw_answer(): """The LLM context must show what the tech actually chose — 'Q? -> yes' is ambiguous for an alternatives question and degrades the next generation.""" ctx = atb._build_context("login issue", "account_login", [ {"node_type": "question", "id": "n1", "text": "Microsoft account or local account?", "answer": "yes", "answer_label": "Microsoft account"}, ]) assert "-> Microsoft account" in ctx assert "-> yes" not in ctx def test_normalize_walked_path_preserves_question_labels(): walked = [ {"node_type": "question", "id": "n1", "text": "Wi-Fi or Ethernet?", "answer": "yes", "answer_label": "Wi-Fi", "yes_label": "Wi-Fi", "no_label": "Ethernet"}, {"node_type": "resolved", "id": "n2", "text": "Fixed."}, ] tree = atb.normalize_walked_path(walked) n1 = tree["nodes"]["n1"] assert n1["yes_label"] == "Wi-Fi" and n1["no_label"] == "Ethernet" def test_validate_node_rejects_hard_floor_text(): node = {"node_type": "instruction", "id": "n1", "text": "Open regedit and change the key", "next": "generate"} with pytest.raises(atb.UnsafeNodeError): atb.validate_node(node) def test_validate_node_accepts_safe_instruction(): node = {"node_type": "instruction", "id": "n1", "text": "Restart the printer.", "next": "generate"} assert atb.validate_node(node)["node_type"] == "instruction" def test_depth_cap_forces_escalate(): walked = [{"node_type": "question", "id": f"n{i}", "text": "?", "answer": "no"} for i in range(atb.MAX_DEPTH)] node = atb.escalate_if_depth_exceeded(walked) assert node is not None and node["node_type"] == "escalate" def test_normalize_walked_path_builds_valid_tree(): walked = [ {"node_type": "question", "id": "n1", "text": "Powered on?", "answer": "no"}, {"node_type": "instruction", "id": "n2", "text": "Power it on.", "answer": "ack"}, {"node_type": "resolved", "id": "n3", "text": "Fixed."}, ] tree = atb.normalize_walked_path(walked) assert isinstance(tree, dict) and tree.get("id") == "n1" # untraversed 'yes' branch of n1 became a needs_review stub assert any(n["node_type"] == "needs_review" for n in tree["nodes"].values()) def test_normalize_walk_ending_on_question_has_no_none_branches(): walked = [ {"node_type": "question", "id": "n1", "text": "Powered on?", "answer": "no"}, ] tree = atb.normalize_walked_path(walked) n1 = tree["nodes"]["n1"] assert n1["yes_next"] is not None and n1["no_next"] is not None # both branches must reference real nodes present in the tree assert n1["yes_next"] in tree["nodes"] and n1["no_next"] in tree["nodes"] def test_normalize_preserves_escalate_reason_category(): walked = [ {"node_type": "question", "id": "n1", "text": "On?", "answer": "no"}, {"node_type": "escalate", "id": "n2", "text": "Beyond L1.", "reason_category": "exhausted_safe_steps"}, ] tree = atb.normalize_walked_path(walked) assert tree["nodes"]["n2"]["reason_category"] == "exhausted_safe_steps" def test_normalize_empty_walk_returns_needs_review_root(): tree = atb.normalize_walked_path([]) assert tree["id"] in tree["nodes"] assert tree["nodes"][tree["id"]]["node_type"] == "needs_review"