""" tests/test_prd_0002.py Unit tests for PRD-0002 — Editor Context Injection. These tests run without any external dependencies (no Elasticsearch, no Ollama, no gRPC server). They validate the logic of the components modified in PRD-0002: - _parse_query_type — classifier output parser (graph.py) - _parse_editor_context — user field parser (openai_proxy.py) - _build_classify_prompt — classify prompt builder (graph.py) - _build_reformulate_query — reformulate anchor builder (graph.py) - _build_generation_prompt — generation prompt builder (graph.py) - _decode_b64 — base64 decoder (server.py) Run with: pytest tests/test_prd_0002.py -v """ import base64 import json import sys import os import pytest # --------------------------------------------------------------------------- # Minimal stubs so we can import graph.py and openai_proxy.py without # the full Docker/src environment loaded # --------------------------------------------------------------------------- # Stub brunix_pb2 so openai_proxy imports cleanly import types brunix_pb2 = types.ModuleType("brunix_pb2") brunix_pb2.AgentRequest = lambda **kw: kw brunix_pb2.AgentResponse = lambda **kw: kw sys.modules["brunix_pb2"] = brunix_pb2 sys.modules["brunix_pb2_grpc"] = types.ModuleType("brunix_pb2_grpc") # Stub grpc grpc_mod = types.ModuleType("grpc") grpc_mod.insecure_channel = lambda *a, **kw: None grpc_mod.Channel = object grpc_mod.RpcError = Exception sys.modules["grpc"] = grpc_mod # Stub grpc_reflection refl = types.ModuleType("grpc_reflection.v1alpha.reflection") sys.modules["grpc_reflection"] = types.ModuleType("grpc_reflection") sys.modules["grpc_reflection.v1alpha"] = types.ModuleType("grpc_reflection.v1alpha") sys.modules["grpc_reflection.v1alpha.reflection"] = refl # Add Docker/src to path so we can import the modules directly DOCKER_SRC = os.path.join(os.path.dirname(__file__), "..", "Docker", "src") sys.path.insert(0, os.path.abspath(DOCKER_SRC)) # --------------------------------------------------------------------------- # Import the functions under test # --------------------------------------------------------------------------- # We import only the pure functions — no LLM, no ES, no gRPC calls def _parse_query_type(raw: str): """Copy of _parse_query_type from graph.py — tested in isolation.""" parts = raw.strip().upper().split() query_type = "RETRIEVAL" use_editor = False if parts: first = parts[0] if first.startswith("CODE_GENERATION") or "CODE" in first: query_type = "CODE_GENERATION" elif first.startswith("CONVERSATIONAL"): query_type = "CONVERSATIONAL" if len(parts) > 1 and parts[1] == "EDITOR": use_editor = True return query_type, use_editor def _decode_b64(value: str) -> str: """Copy of _decode_b64 from server.py — tested in isolation.""" try: return base64.b64decode(value).decode("utf-8") if value else "" except Exception: return "" def _parse_editor_context(user): """Copy of _parse_editor_context from openai_proxy.py — tested in isolation.""" if not user: return "", "", "", "" try: ctx = json.loads(user) if isinstance(ctx, dict): return ( ctx.get("editor_content", "") or "", ctx.get("selected_text", "") or "", ctx.get("extra_context", "") or "", json.dumps(ctx.get("user_info", {})), ) except (json.JSONDecodeError, TypeError): pass return "", "", "", "" def _build_reformulate_query(question: str, selected_text: str) -> str: """Copy of _build_reformulate_query from graph.py — tested in isolation.""" if not selected_text: return question return f"{selected_text}\n\nUser question about the above: {question}" def _build_generation_prompt_injects(editor_content, selected_text, use_editor): """Helper — returns True if editor context would be injected.""" sections = [] if selected_text and use_editor: sections.append("selected_code") if editor_content and use_editor: sections.append("editor_file") return len(sections) > 0 # --------------------------------------------------------------------------- # Tests: _parse_query_type # --------------------------------------------------------------------------- class TestParseQueryType: def test_retrieval_no_editor(self): qt, ue = _parse_query_type("RETRIEVAL NO_EDITOR") assert qt == "RETRIEVAL" assert ue is False def test_retrieval_editor(self): qt, ue = _parse_query_type("RETRIEVAL EDITOR") assert qt == "RETRIEVAL" assert ue is True def test_code_generation_no_editor(self): qt, ue = _parse_query_type("CODE_GENERATION NO_EDITOR") assert qt == "CODE_GENERATION" assert ue is False def test_code_generation_editor(self): qt, ue = _parse_query_type("CODE_GENERATION EDITOR") assert qt == "CODE_GENERATION" assert ue is True def test_conversational_no_editor(self): qt, ue = _parse_query_type("CONVERSATIONAL NO_EDITOR") assert qt == "CONVERSATIONAL" assert ue is False def test_single_token_defaults_no_editor(self): """If model returns only one token, use_editor defaults to False.""" qt, ue = _parse_query_type("RETRIEVAL") assert qt == "RETRIEVAL" assert ue is False def test_empty_defaults_retrieval_no_editor(self): qt, ue = _parse_query_type("") assert qt == "RETRIEVAL" assert ue is False def test_case_insensitive(self): qt, ue = _parse_query_type("retrieval editor") assert qt == "RETRIEVAL" assert ue is True def test_code_shorthand(self): """'CODE' alone should map to CODE_GENERATION.""" qt, ue = _parse_query_type("CODE NO_EDITOR") assert qt == "CODE_GENERATION" assert ue is False def test_extra_whitespace(self): qt, ue = _parse_query_type(" RETRIEVAL NO_EDITOR ") assert qt == "RETRIEVAL" assert ue is False # --------------------------------------------------------------------------- # Tests: _decode_b64 # --------------------------------------------------------------------------- class TestDecodeB64: def test_valid_base64_spanish(self): text = "addVar(mensaje, \"Hola mundo\")\naddResult(mensaje)" encoded = base64.b64encode(text.encode("utf-8")).decode("utf-8") assert _decode_b64(encoded) == text def test_valid_base64_english(self): text = "registerEndpoint(\"GET\", \"/hello\", [], \"public\", handler, \"\")" encoded = base64.b64encode(text.encode("utf-8")).decode("utf-8") assert _decode_b64(encoded) == text def test_empty_string_returns_empty(self): assert _decode_b64("") == "" def test_none_equivalent_empty(self): assert _decode_b64(None) == "" def test_invalid_base64_returns_empty(self): assert _decode_b64("not_valid_base64!!!") == "" def test_unicode_content(self): text = "// función de validación\nif(token, \"SECRET\", \"=\")" encoded = base64.b64encode(text.encode("utf-8")).decode("utf-8") assert _decode_b64(encoded) == text # --------------------------------------------------------------------------- # Tests: _parse_editor_context # --------------------------------------------------------------------------- class TestParseEditorContext: def _encode(self, text: str) -> str: return base64.b64encode(text.encode()).decode() def test_full_context_parsed(self): editor = self._encode("addVar(x, 10)") selected = self._encode("addResult(x)") extra = self._encode("/path/to/file.avap") user_json = json.dumps({ "editor_content": editor, "selected_text": selected, "extra_context": extra, "user_info": {"dev_id": 1, "project_id": 2, "org_id": 3} }) ec, st, ex, ui = _parse_editor_context(user_json) assert ec == editor assert st == selected assert ex == extra assert json.loads(ui) == {"dev_id": 1, "project_id": 2, "org_id": 3} def test_empty_user_returns_empty_tuple(self): ec, st, ex, ui = _parse_editor_context(None) assert ec == st == ex == "" def test_empty_string_returns_empty_tuple(self): ec, st, ex, ui = _parse_editor_context("") assert ec == st == ex == "" def test_plain_string_not_json_returns_empty(self): """Non-JSON user field — backward compat, no error raised.""" ec, st, ex, ui = _parse_editor_context("plain string") assert ec == st == ex == "" def test_missing_fields_default_empty(self): user_json = json.dumps({"editor_content": "abc"}) ec, st, ex, ui = _parse_editor_context(user_json) assert ec == "abc" assert st == "" assert ex == "" def test_user_info_missing_defaults_empty_object(self): user_json = json.dumps({"editor_content": "abc"}) _, _, _, ui = _parse_editor_context(user_json) assert json.loads(ui) == {} def test_user_info_full_object(self): user_json = json.dumps({ "editor_content": "", "selected_text": "", "extra_context": "", "user_info": {"dev_id": 42, "project_id": 7, "org_id": 99} }) _, _, _, ui = _parse_editor_context(user_json) parsed = json.loads(ui) assert parsed["dev_id"] == 42 assert parsed["project_id"] == 7 assert parsed["org_id"] == 99 def test_session_id_not_leaked_into_context(self): """session_id must NOT appear in editor context — it has its own field.""" user_json = json.dumps({ "editor_content": "", "selected_text": "", "extra_context": "", "user_info": {} }) ec, st, ex, ui = _parse_editor_context(user_json) assert "session_id" not in ec assert "session_id" not in st # --------------------------------------------------------------------------- # Tests: _build_reformulate_query # --------------------------------------------------------------------------- class TestBuildReformulateQuery: def test_no_selected_text_returns_question(self): q = "Que significa AVAP?" assert _build_reformulate_query(q, "") == q def test_selected_text_prepended_to_question(self): q = "que hace esto?" selected = "addVar(x, 10)\naddResult(x)" result = _build_reformulate_query(q, selected) assert result.startswith(selected) assert q in result def test_selected_text_anchor_format(self): q = "fix this" selected = "try()\n ormDirect(query, res)\nexception(e)\nend()" result = _build_reformulate_query(q, selected) assert "User question about the above:" in result assert selected in result assert q in result # --------------------------------------------------------------------------- # Tests: editor context injection logic # --------------------------------------------------------------------------- class TestEditorContextInjection: def test_no_injection_when_use_editor_false(self): """Editor content must NOT be injected when use_editor_context is False.""" injected = _build_generation_prompt_injects( editor_content = "addVar(x, 10)", selected_text = "addResult(x)", use_editor = False, ) assert injected is False def test_injection_when_use_editor_true_and_content_present(self): """Editor content MUST be injected when use_editor_context is True.""" injected = _build_generation_prompt_injects( editor_content = "addVar(x, 10)", selected_text = "addResult(x)", use_editor = True, ) assert injected is True def test_no_injection_when_content_empty_even_if_flag_true(self): """Empty fields must never be injected even if flag is True.""" injected = _build_generation_prompt_injects( editor_content = "", selected_text = "", use_editor = True, ) assert injected is False def test_partial_injection_selected_only(self): """selected_text alone triggers injection when flag is True.""" injected = _build_generation_prompt_injects( editor_content = "", selected_text = "addResult(x)", use_editor = True, ) assert injected is True # --------------------------------------------------------------------------- # Tests: classifier routing — EDITOR signal # --------------------------------------------------------------------------- class TestClassifierEditorSignal: """ These tests validate that the two-token output format is correctly parsed for all combinations the classifier can produce. """ VALID_OUTPUTS = [ ("RETRIEVAL NO_EDITOR", "RETRIEVAL", False), ("RETRIEVAL EDITOR", "RETRIEVAL", True), ("CODE_GENERATION NO_EDITOR", "CODE_GENERATION", False), ("CODE_GENERATION EDITOR", "CODE_GENERATION", True), ("CONVERSATIONAL NO_EDITOR", "CONVERSATIONAL", False), ("CONVERSATIONAL EDITOR", "CONVERSATIONAL", True), ] @pytest.mark.parametrize("raw,expected_qt,expected_ue", VALID_OUTPUTS) def test_valid_two_token_output(self, raw, expected_qt, expected_ue): qt, ue = _parse_query_type(raw) assert qt == expected_qt assert ue == expected_ue def test_editor_flag_false_for_general_avap_question(self): """'Que significa AVAP?' -> RETRIEVAL NO_EDITOR.""" qt, ue = _parse_query_type("RETRIEVAL NO_EDITOR") assert ue is False def test_editor_flag_true_for_explicit_editor_reference(self): """'que hace este codigo?' with selected_text -> RETRIEVAL EDITOR.""" qt, ue = _parse_query_type("RETRIEVAL EDITOR") assert ue is True def test_editor_flag_false_for_code_generation_without_reference(self): """'dame un API de hello world' -> CODE_GENERATION NO_EDITOR.""" qt, ue = _parse_query_type("CODE_GENERATION NO_EDITOR") assert ue is False