replypier/backend/tests/repositories/test_guardrail_event.py

"""
Copyright 2024-2026 ChatterMate

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import json
from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock, patch

import pytest

from app.models.guardrail_event import GuardrailEvent
from app.repositories.guardrail_event import (
    GuardrailEventRepository,
    record_guardrail_events,
)


@pytest.fixture
def patched_session(db):
    """record_guardrail_event opens its own SessionLocal — point it at the
    test session."""
    with patch("app.repositories.guardrail_event.SessionLocal") as mock_sl:
        mock_sl.return_value.__enter__.return_value = db
        mock_sl.return_value.__exit__.return_value = False
        yield db


def _add_event(db, rule="injection.frame_tokens", created_at=None, org_id=None):
    event = GuardrailEvent(
        organization_id=org_id,
        surface="widget",
        layer="inbound",
        rule=rule,
        action="counted",
    )
    if created_at is not None:
        event.created_at = created_at
    db.add(event)
    db.commit()
    return event


class TestRecord:
    def test_record_persists_event(self, patched_session):
        record_guardrail_events(
            rules=["injection.override_instructions"],
            surface="widget",
            layer="inbound",
            action="counted",
            score=1.0,
            char_len=42,
        )
        rows = patched_session.query(GuardrailEvent).all()
        assert len(rows) == 1
        assert rows[0].rule == "injection.override_instructions"
        assert rows[0].action == "counted"
        assert rows[0].char_len == 42

    def test_record_never_raises(self):
        with patch(
            "app.repositories.guardrail_event.SessionLocal",
            side_effect=RuntimeError("db down"),
        ):
            # Must swallow: a failed event write cannot affect the reply.
            record_guardrail_events(
                rules=["x"], surface="widget", layer="inbound", action="counted"
            )

    def test_matched_holds_rule_ids_only(self, patched_session):
        # matched is plain JSON while excerpt is encrypted, so the matched
        # SUBSTRING (the visitor's own words) must never land here.
        record_guardrail_events(
            rules=["injection.frame_tokens", "injection.role_hijack"],
            surface="widget",
            layer="inbound",
            action="blocked",
        )
        rows = patched_session.query(GuardrailEvent).all()
        assert len(rows) == 2, "one row per triggered rule"
        for row in rows:
            assert json.dumps(row.matched) == json.dumps(
                ["injection.frame_tokens", "injection.role_hijack"]
            )

    def test_multiple_rules_share_one_session(self, patched_session):
        record_guardrail_events(
            rules=["a", "b", "c"], surface="widget", layer="inbound", action="counted"
        )
        assert patched_session.query(GuardrailEvent).count() == 3

    def test_empty_rules_writes_nothing(self, patched_session):
        record_guardrail_events(
            rules=[], surface="widget", layer="inbound", action="counted"
        )
        assert patched_session.query(GuardrailEvent).count() == 0

    def test_excerpt_nullable(self, patched_session):
        record_guardrail_events(
            rules=["x"], surface="help_center", layer="inbound", action="counted"
        )
        assert patched_session.query(GuardrailEvent).one().excerpt is None


class TestCountsByRule:
    def test_aggregates_per_rule_in_window(self, db):
        now = datetime.now(timezone.utc)
        _add_event(db, rule="injection.frame_tokens", created_at=now)
        _add_event(db, rule="injection.frame_tokens", created_at=now)
        _add_event(db, rule="offtopic.model_refused", created_at=now)
        # Outside the window on both sides.
        _add_event(db, rule="injection.frame_tokens", created_at=now - timedelta(days=10))
        _add_event(db, rule="injection.frame_tokens", created_at=now + timedelta(days=10))

        counts = GuardrailEventRepository(db).counts_by_rule(
            start=now - timedelta(days=1), end=now + timedelta(days=1)
        )
        assert counts == {
            "injection.frame_tokens": 2,
            "offtopic.model_refused": 1,
        }

    def test_org_filter(self, db, test_organization):
        now = datetime.now(timezone.utc)
        _add_event(db, created_at=now, org_id=test_organization.id)
        _add_event(db, created_at=now, org_id=None)
        counts = GuardrailEventRepository(db).counts_by_rule(
            start=now - timedelta(days=1),
            end=now + timedelta(days=1),
            organization_id=test_organization.id,
        )
        assert counts == {"injection.frame_tokens": 1}


class TestReview:
    def test_mark_false_positive(self, db):
        event = _add_event(db)
        assert event.reviewed is False
        updated = GuardrailEventRepository(db).mark_false_positive(event.id)
        assert updated.reviewed is True
        assert updated.false_positive is True

    def test_mark_false_positive_missing_event(self, db):
        import uuid

        assert GuardrailEventRepository(db).mark_false_positive(uuid.uuid4()) is None