mnemosyne/mnemosyne/llm_manager/tests/test_api.py

"""
Tests for LLM Manager DRF API endpoints.
"""

from decimal import Decimal

from django.contrib.auth import get_user_model
from django.test import TestCase
from django.urls import reverse
from rest_framework.test import APIClient

from llm_manager.models import LLMApi, LLMModel, LLMUsage

User = get_user_model()


class LLMApiEndpointTest(TestCase):
    """Tests for the LLM API endpoints."""

    def setUp(self):
        self.user = User.objects.create_user(username="testuser", password="testpass123")
        self.client = APIClient()
        self.client.force_authenticate(user=self.user)
        self.api = LLMApi.objects.create(
            name="Test API",
            api_type="openai",
            base_url="https://api.example.com/v1",
        )

    def test_api_list(self):
        resp = self.client.get(reverse("llm-manager-api:api_list"))
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(len(resp.data), 1)
        self.assertEqual(resp.data[0]["name"], "Test API")

    def test_api_detail(self):
        resp = self.client.get(reverse("llm-manager-api:api_detail", kwargs={"pk": self.api.pk}))
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(resp.data["name"], "Test API")

    def test_api_not_found(self):
        import uuid

        resp = self.client.get(reverse("llm-manager-api:api_detail", kwargs={"pk": uuid.uuid4()}))
        self.assertEqual(resp.status_code, 404)

    def test_requires_auth(self):
        self.client.force_authenticate(user=None)
        resp = self.client.get(reverse("llm-manager-api:api_list"))
        self.assertIn(resp.status_code, [401, 403])


class LLMModelEndpointTest(TestCase):
    """Tests for the LLM Model API endpoints."""

    def setUp(self):
        self.user = User.objects.create_user(username="testuser", password="testpass123")
        self.client = APIClient()
        self.client.force_authenticate(user=self.user)
        self.api = LLMApi.objects.create(
            name="Test API",
            api_type="openai",
            base_url="https://api.example.com/v1",
        )
        self.model = LLMModel.objects.create(
            api=self.api,
            name="gpt-4o",
            model_type="chat",
            context_window=128000,
        )

    def test_model_list(self):
        resp = self.client.get(reverse("llm-manager-api:model_list"))
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(len(resp.data), 1)

    def test_model_list_filter_type(self):
        LLMModel.objects.create(
            api=self.api,
            name="embed-model",
            model_type="embedding",
            context_window=8191,
        )
        resp = self.client.get(reverse("llm-manager-api:model_list") + "?type=embedding")
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(len(resp.data), 1)
        self.assertEqual(resp.data[0]["model_type"], "embedding")

    def test_model_detail(self):
        resp = self.client.get(reverse("llm-manager-api:model_detail", kwargs={"pk": self.model.pk}))
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(resp.data["name"], "gpt-4o")

    def test_system_models_empty(self):
        resp = self.client.get(reverse("llm-manager-api:system_models"))
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(resp.data, {})

    def test_system_models_configured(self):
        self.model.is_system_chat_model = True
        self.model.save()
        resp = self.client.get(reverse("llm-manager-api:system_models"))
        self.assertEqual(resp.status_code, 200)
        self.assertIn("chat", resp.data)
        self.assertEqual(resp.data["chat"]["name"], "gpt-4o")


class LLMUsageEndpointTest(TestCase):
    """Tests for the LLM Usage API endpoints."""

    def setUp(self):
        self.user = User.objects.create_user(username="testuser", password="testpass123")
        self.client = APIClient()
        self.client.force_authenticate(user=self.user)
        self.api = LLMApi.objects.create(
            name="Test API",
            api_type="openai",
            base_url="https://api.example.com/v1",
        )
        self.model = LLMModel.objects.create(
            api=self.api,
            name="gpt-4o",
            model_type="chat",
            context_window=128000,
            input_cost_per_1k=Decimal("0.0025"),
            output_cost_per_1k=Decimal("0.01"),
        )

    def test_usage_list_empty(self):
        resp = self.client.get(reverse("llm-manager-api:usage_list"))
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(resp.data, [])

    def test_usage_create(self):
        resp = self.client.post(
            reverse("llm-manager-api:usage_list"),
            {
                "model": str(self.model.pk),
                "input_tokens": 1000,
                "output_tokens": 500,
                "purpose": "other",
            },
            format="json",
        )
        self.assertEqual(resp.status_code, 201)
        self.assertEqual(LLMUsage.objects.count(), 1)

    def test_usage_list_returns_own_records(self):
        other_user = User.objects.create_user(username="other", password="testpass123")
        LLMUsage.objects.create(user=self.user, model=self.model, input_tokens=100, output_tokens=50)
        LLMUsage.objects.create(user=other_user, model=self.model, input_tokens=200, output_tokens=100)
        resp = self.client.get(reverse("llm-manager-api:usage_list"))
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(len(resp.data), 1)