GIA/core/tests/test_command_variant_policy.py

from __future__ import annotations

from unittest.mock import AsyncMock, patch

from asgiref.sync import async_to_sync
from django.test import TransactionTestCase

from core.commands.base import CommandContext
from core.commands.handlers.bp import BPCommandHandler
from core.commands.policies import ensure_variant_policies_for_profile
from core.models import (
    BusinessPlanDocument,
    AI,
    ChatSession,
    CommandAction,
    CommandChannelBinding,
    CommandProfile,
    CommandVariantPolicy,
    Message,
    Person,
    PersonIdentifier,
    User,
)


class CommandVariantPolicyTests(TransactionTestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username="variant-user",
            email="variant@example.com",
            password="x",
        )
        self.person = Person.objects.create(user=self.user, name="Variant Person")
        self.identifier = PersonIdentifier.objects.create(
            user=self.user,
            person=self.person,
            service="whatsapp",
            identifier="120363402761690215",
        )
        self.session = ChatSession.objects.create(user=self.user, identifier=self.identifier)
        self.profile = CommandProfile.objects.create(
            user=self.user,
            slug="bp",
            name="Business Plan",
            enabled=True,
            trigger_token="#bp#",
            reply_required=True,
            exact_match_only=True,
            visibility_mode="status_in_source",
            template_text="TEMPLATE SHOULD NOT LEAK INTO bp set",
        )
        AI.objects.create(
            user=self.user,
            base_url="https://example.invalid",
            api_key="test-key",
            model="gpt-4o-mini",
        )
        CommandAction.objects.create(
            profile=self.profile,
            action_type="extract_bp",
            enabled=True,
            position=0,
        )
        CommandAction.objects.create(
            profile=self.profile,
            action_type="save_document",
            enabled=True,
            position=1,
        )
        CommandAction.objects.create(
            profile=self.profile,
            action_type="post_result",
            enabled=True,
            position=2,
        )
        CommandChannelBinding.objects.create(
            profile=self.profile,
            direction="ingress",
            service="whatsapp",
            channel_identifier="120363402761690215",
            enabled=True,
        )
        CommandChannelBinding.objects.create(
            profile=self.profile,
            direction="egress",
            service="whatsapp",
            channel_identifier="120363402761690215",
            enabled=True,
        )

    def _ctx(self, trigger: Message, text: str) -> CommandContext:
        return CommandContext(
            service="whatsapp",
            channel_identifier="120363402761690215",
            message_id=str(trigger.id),
            user_id=self.user.id,
            message_text=text,
            payload={},
        )

    def test_ensure_variant_policies_backfills_bp_defaults(self):
        rows = ensure_variant_policies_for_profile(self.profile)
        self.assertSetEqual(set(rows.keys()), {"bp", "bp_set", "bp_set_range"})
        self.assertEqual("ai", rows["bp"].generation_mode)
        self.assertEqual("verbatim", rows["bp_set"].generation_mode)
        self.assertEqual("verbatim", rows["bp_set_range"].generation_mode)
        self.assertTrue(rows["bp"].send_plan_to_egress)
        self.assertTrue(rows["bp"].send_status_to_source)

    def test_bp_primary_can_run_in_verbatim_mode_without_ai(self):
        ensure_variant_policies_for_profile(self.profile)
        policy = CommandVariantPolicy.objects.get(profile=self.profile, variant_key="bp")
        policy.generation_mode = "verbatim"
        policy.send_plan_to_egress = False
        policy.send_status_to_source = False
        policy.send_status_to_egress = False
        policy.save()

        anchor = Message.objects.create(
            user=self.user,
            session=self.session,
            sender_uuid="peer",
            text="anchor line",
            ts=1000,
            source_service="whatsapp",
            source_chat_id="120363402761690215",
        )
        trigger = Message.objects.create(
            user=self.user,
            session=self.session,
            sender_uuid="me",
            text="#bp#",
            ts=2000,
            source_service="whatsapp",
            source_chat_id="120363402761690215",
            reply_to=anchor,
        )

        result = async_to_sync(BPCommandHandler().execute)(self._ctx(trigger, "#bp#"))
        self.assertTrue(result.ok)
        doc = BusinessPlanDocument.objects.get(trigger_message=trigger)
        self.assertEqual("anchor line\n#bp#", doc.content_markdown)

    def test_bp_set_ai_mode_ignores_template(self):
        ensure_variant_policies_for_profile(self.profile)
        policy = CommandVariantPolicy.objects.get(profile=self.profile, variant_key="bp_set")
        policy.generation_mode = "ai"
        policy.send_plan_to_egress = False
        policy.send_status_to_source = False
        policy.send_status_to_egress = False
        policy.save()

        trigger = Message.objects.create(
            user=self.user,
            session=self.session,
            sender_uuid="me",
            text="#bp set# text to transform",
            ts=1000,
            source_service="whatsapp",
            source_chat_id="120363402761690215",
        )

        with patch(
            "core.commands.handlers.bp.ai_runner.run_prompt",
            new=AsyncMock(return_value="AI RESULT"),
        ) as mocked:
            result = async_to_sync(BPCommandHandler().execute)(
                self._ctx(trigger, trigger.text)
            )

        self.assertTrue(result.ok)
        doc = BusinessPlanDocument.objects.get(trigger_message=trigger)
        self.assertEqual("AI RESULT", doc.content_markdown)
        call_args = mocked.await_args.args
        prompt_payload = call_args[0]
        self.assertNotIn("TEMPLATE SHOULD NOT LEAK", str(prompt_payload))

    def test_delivery_flags_control_source_and_egress_status(self):
        ensure_variant_policies_for_profile(self.profile)
        policy = CommandVariantPolicy.objects.get(
            profile=self.profile,
            variant_key="bp_set_range",
        )
        policy.generation_mode = "verbatim"
        policy.store_document = False
        policy.send_plan_to_egress = False
        policy.send_status_to_source = True
        policy.send_status_to_egress = True
        policy.save()

        anchor = Message.objects.create(
            user=self.user,
            session=self.session,
            sender_uuid="peer",
            text="line one",
            ts=1000,
            source_service="whatsapp",
            source_chat_id="120363402761690215",
        )
        trigger = Message.objects.create(
            user=self.user,
            session=self.session,
            sender_uuid="me",
            text="#bp set range#",
            ts=2000,
            source_service="whatsapp",
            source_chat_id="120363402761690215",
            reply_to=anchor,
        )

        with patch(
            "core.commands.handlers.bp.post_status_in_source",
            new=AsyncMock(return_value=True),
        ) as source_status, patch(
            "core.commands.handlers.bp.post_to_channel_binding",
            new=AsyncMock(return_value=True),
        ) as binding_send:
            result = async_to_sync(BPCommandHandler().execute)(
                self._ctx(trigger, trigger.text)
            )

        self.assertTrue(result.ok)
        source_status.assert_awaited()
        self.assertEqual(1, binding_send.await_count)
        self.assertFalse(BusinessPlanDocument.objects.filter(trigger_message=trigger).exists())