fa-agent/tests/test_assistant_api.py

import json
from unittest.mock import patch

import httpx
import pytest
from django.contrib.auth.models import User
from django.test import Client

from finance.assistant import model_intent
from finance.models import AIConfig, Audit, FormulaDefinition, FormulaExport, Membership

pytestmark = pytest.mark.django_db


def completed_response(intent=None):
    class Response:
        def raise_for_status(self):
            return None

        def json(self):
            return {
                "status": "completed",
                "output": [
                    {
                        "content": [
                            {
                                "type": "output_text",
                                "text": json.dumps(
                                    intent
                                    or {
                                        "action": "overview",
                                        "metric": "revenue",
                                        "group": "department",
                                        "filters": [{"key": "department", "value": "儿科"}],
                                    }
                                ),
                            }
                        ]
                    }
                ],
            }

    return Response()


def test_ai_config_requires_full_admin_and_masks_key(workspace):
    ds, member, client, _, _ = workspace
    url = "/api/ai-config"

    viewer = User.objects.create_user("ai-viewer", password="Finance-TestPassword991!")
    Membership.objects.create(user=viewer, organization=ds.organization, role="viewer")
    viewer_client = Client()
    assert viewer_client.login(username="ai-viewer", password="Finance-TestPassword991!")
    assert viewer_client.get(url).status_code == 403

    member.departments = ["儿科"]
    member.save(update_fields=["departments"])
    scoped = Client()
    assert scoped.login(username=member.user.username, password="Finance-TestPassword991!")
    assert scoped.get(url).status_code == 403
    member.departments = []
    member.save(update_fields=["departments"])

    assert client.get(url).json() == {
        "enabled": False,
        "provider": "local",
        "base_url": "",
        "model": "",
        "api_key_set": False,
        "updated_at": None,
    }

    invalid = [
        {"enabled": True, "provider": "local", "base_url": "", "model": "qwen2.5:7b"},
        {"enabled": True, "provider": "openai", "base_url": "", "model": "gpt-4o"},
        {"enabled": True, "provider": "local", "base_url": "ftp://x", "model": "m"},
        {"enabled": True, "provider": "local", "base_url": "http://x/v1", "model": ""},
    ]
    for payload in invalid:
        assert client.put(url, data=payload, content_type="application/json").status_code == 400

    saved = client.put(
        url,
        data={
            "enabled": True,
            "provider": "local",
            "base_url": "http://127.0.0.1:11434/v1",
            "api_key": "ollama-secret",
            "model": "qwen2.5:7b",
        },
        content_type="application/json",
    )
    assert saved.status_code == 200
    shown = client.get(url).json()
    assert shown["api_key_set"] is True
    assert shown["enabled"] is True and shown["provider"] == "local"
    assert "ollama-secret" not in json.dumps(shown)
    assert client.get("/api/session").json()["ai_enabled"] is True
    assert Audit.objects.filter(action="ai_config.update").exists()


def test_ai_models_lists_local_models(workspace):
    ds, _, client, _, _ = workspace
    client.put(
        "/api/ai-config",
        data={
            "enabled": False,
            "provider": "local",
            "base_url": "http://127.0.0.1:11434/v1",
            "model": "",
        },
        content_type="application/json",
    )
    with patch("finance.ai.httpx.Client") as http:
        http.return_value.__enter__.return_value.get.return_value = type(
            "Response",
            (),
            {
                "raise_for_status": lambda self: None,
                "json": lambda self: {"data": [{"id": "qwen2.5:7b"}, {"id": "llama3.1:8b"}]},
            },
        )()
        listed = client.get("/api/ai-models")
    assert listed.status_code == 200
    assert listed.json()["items"] == ["qwen2.5:7b", "llama3.1:8b"]
    call = http.return_value.__enter__.return_value.get.call_args
    assert call.args[0] == "http://127.0.0.1:11434/v1/models"

    with patch("finance.ai.httpx.Client") as http:
        http.return_value.__enter__.return_value.get.side_effect = httpx.ConnectError("boom")
        failed = client.get("/api/ai-models")
    assert failed.status_code == 502


def test_assistant_uses_org_db_config(workspace):
    ds, _, client, *_ = workspace
    client.put(
        "/api/ai-config",
        data={
            "enabled": True,
            "provider": "local",
            "base_url": "http://127.0.0.1:11434/v1",
            "model": "qwen2.5:7b",
        },
        content_type="application/json",
    )
    with patch("finance.assistant.httpx.Client") as http:
        http.return_value.__enter__.return_value.post.return_value = completed_response()
        response = client.post(
            f"/api/datasets/{ds.pk}/assistant",
            data={
                "question": "分析儿科收入",
                "context": {"start": "2026-08", "end": "2026-08", "comparison": "yoy", "group": "department"},
                "use_model": True,
                "consent": True,
            },
            content_type="application/json",
        )
    assert response.status_code == 200
    assert response.json()["provider"] == "local"
    call = http.return_value.__enter__.return_value.post.call_args
    assert call.args[0] == "http://127.0.0.1:11434/v1/responses"
    assert call.kwargs["json"]["model"] == "qwen2.5:7b"
    assert "Authorization" not in call.kwargs["headers"]

    disabled = AIConfig.objects.get(organization=ds.organization)
    disabled.enabled = False
    disabled.save(update_fields=["enabled"])
    denied = client.post(
        f"/api/datasets/{ds.pk}/assistant",
        data={"question": "分析儿科收入", "use_model": True, "consent": True},
        content_type="application/json",
    )
    assert denied.status_code == 503


def test_model_intent_accepts_only_valid_dimensions_and_sends_no_rows(workspace, settings):
    ds, _, _, _, _ = workspace
    settings.AI_ENABLED = True
    settings.OPENAI_MODEL = "test-model"
    settings.OPENAI_API_KEY = "test-key"

    class Response:
        def raise_for_status(self):
            return None

        def json(self):
            return {
                "status": "completed",
                "output": [
                    {
                        "content": [
                            {
                                "type": "output_text",
                                "text": json.dumps(
                                    {
                                        "action": "overview",
                                        "metric": "revenue",
                                        "group": "department",
                                        "filters": [{"key": "department", "value": "儿科"}],
                                    }
                                ),
                            }
                        ]
                    }
                ],
            }

    with patch("finance.assistant.httpx.Client") as http:
        http.return_value.__enter__.return_value.post.return_value = Response()
        intent = model_intent("看儿科收入", ds, {"metric": "revenue", "group": "department", "filters": {}})

    assert intent == {
        "action": "overview",
        "metric": "revenue",
        "group": "department",
        "filters": {"department": "儿科"},
    }
    payload = http.return_value.__enter__.return_value.post.call_args.kwargs["json"]
    assert payload["store"] is False
    assert "C1" not in json.dumps(payload)


def test_model_intent_rejects_unknown_filter_and_incomplete_provider_response(workspace, settings):
    ds, *_ = workspace
    settings.AI_ENABLED = True
    settings.OPENAI_MODEL = "test-model"
    settings.OPENAI_API_KEY = "test-key"

    class Response:
        def raise_for_status(self):
            return None

        def json(self):
            return {
                "status": "completed",
                "output": [
                    {
                        "content": [
                            {
                                "type": "output_text",
                                "text": '{"action":"overview","metric":"revenue","group":"department","filters":[{"key":"__class__","value":"x"}]}',
                            }
                        ]
                    }
                ],
            }

    response = Response()
    with patch("finance.assistant.httpx.Client") as http:
        http.return_value.__enter__.return_value.post.return_value = response
        with pytest.raises(Exception, match="无效意图"):
            model_intent("分析收入", ds, {"metric": "revenue", "group": "department", "filters": {}})

    with patch("finance.assistant.httpx.Client") as http:
        http.return_value.__enter__.return_value.post.return_value = type(
            "Response",
            (),
            {"raise_for_status": lambda self: None, "json": lambda self: {"status": "in_progress"}},
        )()
        with pytest.raises(Exception, match="无效意图"):
            model_intent("分析收入", ds, {"metric": "revenue", "group": "department", "filters": {}})


def test_assistant_local_response_and_external_consent_guard(workspace, settings):
    ds, _, client, *_ = workspace
    url = f"/api/datasets/{ds.pk}/assistant"
    payload = {
        "question": "查看实际收入",
        "context": {"start": "2026-08", "end": "2026-08", "comparison": "yoy", "group": "department"},
    }
    response = client.post(url, data=payload, content_type="application/json")
    assert response.status_code == 200
    assert response.json()["provider"] == "local"
    assert Audit.objects.filter(action="assistant.query").exists()

    settings.AI_ENABLED = True
    settings.OPENAI_API_KEY = "test-key"
    settings.OPENAI_MODEL = "test-model"
    denied = client.post(
        url,
        data={**payload, "use_model": True, "consent": False},
        content_type="application/json",
    )
    assert denied.status_code == 400
    assert "同意" in denied.json()["error"]


def test_formula_design_and_export_api_validate_consent_and_scope(workspace, settings):
    ds, _, client, *_ = workspace
    url = f"/api/datasets/{ds.pk}"
    payload = {
        "context": {"start": "2026-08", "end": "2026-08", "comparison": "yoy", "group": "department"},
        "messages": [{"role": "user", "content": "收入成本差"}],
        "consent": False,
    }
    denied = client.post(f"{url}/formula-design", data=payload, content_type="application/json")
    assert denied.status_code == 400

    export = client.post(
        f"{url}/formula-exports",
        data={**payload["context"], "custom_columns": [{"expression": {}}]},
        content_type="application/json",
    )
    assert export.status_code == 400
    assert FormulaExport.objects.count() == 0

    scoped = Client()
    user = ds.created_by
    user.membership.departments = ["儿科"]
    user.membership.save(update_fields=["departments"])
    assert scoped.login(username=user.username, password="Finance-TestPassword991!")
    forbidden = scoped.post(
        f"{url}/formula-exports", data=payload["context"], content_type="application/json"
    )
    assert forbidden.status_code == 403
    assert scoped.get(f"{url}/formula-definitions").status_code == 403


def test_formula_definition_library_is_private_validated_and_audited(workspace):
    ds, _, client, _, _ = workspace
    url = f"/api/datasets/{ds.pk}/formula-definitions"
    definition = {
        "name": "收入成本差",
        "meaning": "实际收入减成本",
        "scope": "所选分组与期间",
        "missing": "结果不可用（#N/A）；零分母不可用",
        "expression": {
            "operator": "subtract",
            "left": {"metric": "revenue"},
            "right": {"metric": "cost"},
        },
    }
    saved = client.post(url, data={"definition": definition}, content_type="application/json")
    assert saved.status_code == 201
    definition_id = saved.json()["id"]
    assert client.get(url).json()["items"][0]["definition"] == definition

    updated = client.post(
        url,
        data={"definition": {**definition, "name": "  收入成本差  ", "meaning": "更新说明"}},
        content_type="application/json",
    )
    assert updated.status_code == 200
    assert FormulaDefinition.objects.filter(dataset=ds, user=ds.created_by).count() == 1
    assert updated.json()["definition"]["meaning"] == "更新说明"

    other = Client()
    other_user = ds.created_by.__class__.objects.create_user(
        "formula-library-other", password="Formula-Library-TestPassword991!"
    )
    Membership.objects.create(user=other_user, organization=ds.organization, role="analyst")
    assert other.login(username=other_user.username, password="Formula-Library-TestPassword991!")
    assert other.get(url).json()["items"] == []
    assert other.delete(f"{url}/{definition_id}").status_code == 404

    deleted = client.delete(f"{url}/{definition_id}")
    assert deleted.status_code == 200
    assert not FormulaDefinition.objects.filter(pk=definition_id).exists()
    assert Audit.objects.filter(action="formula_definition.save").count() == 2
    assert Audit.objects.filter(action="formula_definition.delete").count() == 1


def test_formula_definition_library_rejects_unavailable_metric(workspace):
    ds, _, client, *_ = workspace
    mapping = {**ds.mapping, "columns": [c for c in ds.mapping["columns"] if c["role"] != "cost"]}
    ds.mapping = mapping
    ds.save(update_fields=["mapping"])
    response = client.post(
        f"/api/datasets/{ds.pk}/formula-definitions",
        data={
            "definition": {
                "name": "收入成本差",
                "meaning": "实际收入减成本",
                "scope": "所选分组与期间",
                "missing": "结果不可用（#N/A）；零分母不可用",
                "expression": {
                    "operator": "subtract",
                    "left": {"metric": "revenue"},
                    "right": {"metric": "cost"},
                },
            }
        },
        content_type="application/json",
    )
    assert response.status_code == 400
    assert "成本" in response.json()["error"]