stack/tests/rex/comments/test_classify.py

"""Tests for rex.comments.classify — pure logic, no network.

The LLM transport is exercised with a fake Provider; prompt/schema
contents are asserted structurally so vocabulary drift breaks loudly.
"""

from __future__ import annotations

import pytest

from prisma.llm import LLMResult
from rex.comments.classify import (
    CLASSIFY_TOOL,
    POSITIONS,
    PROVISIONS,
    STAKEHOLDERS,
    THEMES,
    build_call,
    classify,
    is_relevant,
    strip_html,
    validate,
)


class TestRelevance:
    @pytest.mark.parametrize(
        "text",
        [
            "We oppose the skin substitute payment change.",
            "Skin Substitutes should stay at ASP+6%",
            "cellular and tissue-based products (CTPs)",
            "the new C5271 application code",
            "HCPCS Q4186 pricing",
        ],
    )
    def test_relevant(self, text):
        assert is_relevant(text)

    @pytest.mark.parametrize(
        "text",
        [
            "I am concerned about telehealth reimbursement.",
            "OT students deserve better PFS treatment",
            "",
        ],
    )
    def test_irrelevant(self, text):
        assert not is_relevant(text)


class TestStripHtml:
    def test_strips_tags_and_entities(self):
        raw = '<div class="zotero-note"><h2>Inline</h2><p>ASP&amp;6% &lt;rate&gt;</p></div>'
        assert strip_html(raw) == "Inline ASP&6% <rate>"


class TestValidate:
    def _raw(self, **over):
        base = {
            "position": "oppose",
            "themes": ["patient_access", "asp_methodology"],
            "stakeholder_type": "provider",
            "provisions": ["flat_rate"],
            "commenter_name": "Dr. A",
            "organization": "Clinic B",
            "rationale": "cites access",
        }
        base.update(over)
        return base

    def test_happy_path(self):
        rec = validate(self._raw())
        assert rec["position_score"] == -1
        assert rec["themes"] == "patient_access; asp_methodology"
        assert rec["provisions"] == "flat_rate"

    def test_invalid_position_rejected(self):
        with pytest.raises(ValueError, match="invalid position"):
            validate(self._raw(position="meh"))

    def test_unknown_enum_values_dropped(self):
        rec = validate(
            self._raw(
                themes=["patient_access", "made_up"],
                provisions=["flat_rate", "nonsense"],
                stakeholder_type="alien",
            )
        )
        assert rec["themes"] == "patient_access"
        assert rec["provisions"] == "flat_rate"
        assert rec["stakeholder_type"] == "unknown"

    def test_dedupes_preserving_order(self):
        rec = validate(self._raw(themes=["flat_rate_design", "flat_rate_design"]))
        assert rec["themes"] == "flat_rate_design"


class TestBuildCall:
    def test_shape(self):
        call = build_call("some comment text", title="CMS-2025-0304-0004")
        assert call.force_tool == "classify_comment"
        assert call.tools == [CLASSIFY_TOOL]
        assert call.messages[0].role == "system"
        assert call.messages[0].cache is True
        assert "CMS-2025-0304-0004" in call.messages[1].content

    def test_truncates_long_text(self):
        call = build_call("x" * 100_000)
        assert len(call.messages[1].content) <= 30_000

    def test_schema_vocab_in_sync(self):
        props = CLASSIFY_TOOL.schema["properties"]
        assert props["position"]["enum"] == list(POSITIONS)
        assert props["themes"]["items"]["enum"] == list(THEMES)
        assert props["stakeholder_type"]["enum"] == list(STAKEHOLDERS)
        assert props["provisions"]["items"]["enum"] == list(PROVISIONS)


class _FakeProvider:
    def __init__(self, tool_calls):
        self._tool_calls = tool_calls
        self.calls = []

    def complete(self, call):
        self.calls.append(call)
        return LLMResult(text="", tool_calls=self._tool_calls, usage={})


class TestClassify:
    def test_returns_validated_record(self):
        provider = _FakeProvider(
            [
                {
                    "name": "classify_comment",
                    "input": {
                        "position": "strongly_support",
                        "themes": ["fraud_waste_abuse"],
                        "stakeholder_type": "government",
                        "provisions": ["reclassification"],
                    },
                }
            ]
        )
        rec = classify(provider, "text", title="C-1")
        assert rec["position_score"] == 2
        assert provider.calls[0].force_tool == "classify_comment"

    def test_no_tool_call_raises(self):
        provider = _FakeProvider([])
        with pytest.raises(ValueError, match="no classify_comment tool call"):
            classify(provider, "text")


def test_docket_rules_maps_cy2026_pfs():
    from rex.comments.classify import DOCKET_RULES

    assert DOCKET_RULES["CMS-2025-0304"] == "CMS-1832-P"


def test_docket_rules_maps_cy2027_pfs():
    from rex.comments.classify import DOCKET_RULES

    assert DOCKET_RULES["CMS-2026-2377"] == "CMS-1848-P"


def test_system_prompt_names_the_pfs_rule():
    from rex.comments.classify import _SYSTEM

    assert "CMS-1832-P" in _SYSTEM
    assert "CMS-1834-P" not in _SYSTEM
    assert "PFS proposed rule" in _SYSTEM