vignette/apps/api/app/test_orchestrator_masking.py

"""Regression tests for P1 PII masking before engine requests."""

from __future__ import annotations

import json
import unittest
from typing import Any
from unittest.mock import patch

from .engine_client import EngineClient, GenerateResponse
from .services import guardrail, orchestrator, persona, state_machine


RAW_PHONE = "010-1234-5678"
RAW_EMAIL = "test@example.com"
RAW_RRN = "990101-1234567"
RAW_TEXT = f"My phone is {RAW_PHONE}, email {RAW_EMAIL}, and RRN {RAW_RRN}."
RAW_VALUES = (RAW_PHONE, RAW_EMAIL, RAW_RRN)
MASK_VALUES = ("[PHONE]", "[EMAIL]", "[RRN]")


def _json_blob(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, default=str)


def _message_blob(messages: object) -> str:
    return "\n".join(message.content for message in messages)  # type: ignore[attr-defined]


def _initial_state() -> state_machine.SessionState:
    card = persona.P1
    return state_machine.init_state(
        params=card.openness_params(),
    )


def _prepare_context() -> orchestrator.TurnContext:
    return orchestrator.prepare_turn(
        session_id="masking-session",
        case_id="masking-case",
        card=persona.P1,
        state=_initial_state(),
        learner_text=RAW_TEXT,
        recent_turns=[
            {
                "speaker": "counselor",
                "text": "Previous learner contact was already masked: [PHONE] [EMAIL] [RRN].",
            }
        ],
    )


def _assert_no_raw_pii(test: unittest.TestCase, value: object) -> None:
    blob = _json_blob(value)
    for raw in RAW_VALUES:
        test.assertNotIn(raw, blob)


def _assert_masked_pii_present(test: unittest.TestCase, value: object) -> None:
    blob = _json_blob(value)
    for masked in MASK_VALUES:
        test.assertIn(masked, blob)


class CaptureGenerateEngine:
    def __init__(self) -> None:
        self.request = None
        self.payload: dict[str, Any] | None = None
        self._payload_builder = EngineClient(base_url="http://engine.test")

    async def generate(self, req):
        self.request = req
        self.payload = self._payload_builder._payload(req)
        return GenerateResponse(
            text="Masked engine reply.",
            model="fake-model",
            provider="fake-provider",
            tokens_in=3,
            tokens_out=4,
            cost_usd=0.0,
        )


class CaptureStreamEngine:
    engine_mode = "fake-provider"
    default_model = "fake-model"

    def __init__(self) -> None:
        self.request = None
        self.payload: dict[str, Any] | None = None
        self._payload_builder = EngineClient(base_url="http://engine.test")

    async def stream(self, req):
        self.request = req
        self.payload = self._payload_builder._payload(req)
        yield "event: token"
        yield 'data: {"text":"Masked stream reply."}'
        yield "event: done"
        yield (
            'data: {"provider":"fake-provider","model":"fake-model",'
            '"tokens_in":5,"tokens_out":6,"cost_usd":0.0}'
        )


class OrchestratorMaskingGateTest(unittest.IsolatedAsyncioTestCase):
    def setUp(self) -> None:
        self.presidio_patch = patch.object(
            guardrail,
            "_try_load_presidio",
            return_value=(None, None),
        )
        self.presidio_patch.start()
        self.addCleanup(self.presidio_patch.stop)

    def test_prepare_turn_keeps_raw_text_but_builds_masked_engine_messages(self) -> None:
        ctx = _prepare_context()

        self.assertEqual(ctx.learner_text_raw, RAW_TEXT)
        for raw in RAW_VALUES:
            self.assertIn(raw, ctx.learner_text_raw)
            self.assertNotIn(raw, ctx.learner_text_masked)
            self.assertNotIn(raw, _message_blob(ctx.messages))

        for masked in MASK_VALUES:
            self.assertIn(masked, ctx.learner_text_masked)
            self.assertIn(masked, _message_blob(ctx.messages))

    def test_prepare_turn_masks_raw_pii_from_context_inputs(self) -> None:
        ctx = orchestrator.prepare_turn(
            session_id="masking-session",
            case_id="masking-case",
            card=persona.P1,
            state=_initial_state(),
            learner_text="Current text has no identifiers.",
            recall_summary=f"Recall mentioned {RAW_PHONE}.",
            pinned_facts=[f"Pinned email {RAW_EMAIL}."],
            recent_turns=[
                {"speaker": "counselor", "text": f"Previous raw RRN {RAW_RRN}."},
            ],
        )

        blob = _message_blob(ctx.messages)
        for raw in RAW_VALUES:
            self.assertNotIn(raw, blob)
        for masked in MASK_VALUES:
            self.assertIn(masked, blob)

    async def test_run_turn_generate_sends_only_masked_engine_payload(self) -> None:
        ctx = _prepare_context()
        engine = CaptureGenerateEngine()

        await orchestrator.run_turn_generate(ctx, engine)  # type: ignore[arg-type]

        self.assertIsNotNone(engine.request)
        self.assertIsNotNone(engine.payload)
        _assert_no_raw_pii(self, engine.request.messages)
        _assert_no_raw_pii(self, engine.payload)
        _assert_masked_pii_present(self, engine.request.messages)
        _assert_masked_pii_present(self, engine.payload)
        self.assertEqual(ctx.learner_text_raw, RAW_TEXT)

    async def test_run_turn_stream_sends_only_masked_engine_payload(self) -> None:
        ctx = _prepare_context()
        engine = CaptureStreamEngine()

        events = [
            event
            async for event in orchestrator.run_turn_stream(ctx, engine)  # type: ignore[arg-type]
        ]

        self.assertEqual([event.event for event in events], ["token", "done"])
        self.assertIsNotNone(engine.request)
        self.assertIsNotNone(engine.payload)
        _assert_no_raw_pii(self, engine.request.messages)
        _assert_no_raw_pii(self, engine.payload)
        _assert_masked_pii_present(self, engine.request.messages)
        _assert_masked_pii_present(self, engine.payload)
        self.assertEqual(ctx.learner_text_raw, RAW_TEXT)


if __name__ == "__main__":
    unittest.main()