Testowanie kodu z API LLM: pytest i unittest.mock krok po kroku

Jeśli test Twojego kodu wysyła prawdziwe zapytanie do API, ma trzy wady: kosztuje, jest wolny i za każdym razem dostaje inną odpowiedź. Do tego wymaga klucza API, którego nie chcesz trzymać w środowisku CI. W tym artykule pokażemy, jak napisać testy z atrapami, żeby sprawdzać własny kod, a nie model.

To druga część cyklu. W pierwszej napisaliśmy kod do wywołań, tokenów i kosztów, a w trzeciej zbudujemy z niego projekt.

Dlaczego nie warto testować na prawdziwym API?

Testujemy to, co należy do nas: czy poprawnie czytamy odpowiedź, czy liczymy koszt, czy ponawiamy zapytanie po błędzie. Jakość odpowiedzi modelu to osobny temat i nie da się jej sprawdzić zwykłym assert.

Testy to też najlepsza ochrona przy vibe codingu: kod wygenerowany przez AI działa tylko do czasu, aż ktoś sprawdzi go testem.

Kod, który będziemy testować

Użyjemy funkcji ask_model, która wysyła zapytanie, zlicza zużycie tokenów i ponawia próbę po błędzie limitu:

import time

from openai import APITimeoutError, RateLimitError

MODEL = "gpt-4o-mini"


def ask_model(client, prompt, report, retries=3, sleep=time.sleep):
    for attempt in range(retries + 1):
        try:
            response = client.chat.completions.create(
                model=MODEL,
                messages=[{"role": "user", "content": prompt}],
                timeout=30,
            )
        except (RateLimitError, APITimeoutError):
            if attempt == retries:
                raise
            sleep(2 ** attempt)
            continue
        report.add(response.usage)
        return response.choices[0].message.content

Zwróć uwagę na dwie decyzje projektowe, które ułatwiają testowanie. Klient (client) jest przekazywany jako argument, więc w teście podamy atrapę. Funkcja sleep też jest argumentem, więc test nie czeka sekundami na kolejną próbę.

Jak zrobić atrapę klienta w unittest.mock?

MagicMock to obiekt, który przyjmie dowolne wywołanie i dowolny atrybut. Wystarczy powiedzieć, co ma zwrócić:

from types import SimpleNamespace
from unittest.mock import MagicMock


def fake_response(text="Streszczenie.", prompt_tokens=100, completion_tokens=20):
    return SimpleNamespace(
        choices=[SimpleNamespace(message=SimpleNamespace(content=text))],
        usage=SimpleNamespace(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens),
    )


def test_ask_model_returns_text_and_counts_usage():
    client = MagicMock()
    client.chat.completions.create.return_value = fake_response("Cześć", 50, 10)
    report = UsageReport()

    assert ask_model(client, "Hej", report) == "Cześć"
    assert (report.calls, report.prompt_tokens, report.completion_tokens) == (1, 50, 10)

Uruchamiasz to poleceniem pytest. Prosty obiekt SimpleNamespace wystarcza za odpowiedź, bo nasz kod czyta tylko choices, message.content i usage. Nie musisz odtwarzać całej struktury z biblioteki.

Jak współdzielić atrapę między testami z fixture?

Jeśli ten sam klient potrzebny jest w wielu testach, wyciągnij go do fixture w pytest:

import pytest


@pytest.fixture
def client():
    return MagicMock()


def test_ask_model_uses_fixture(client):
    client.chat.completions.create.return_value = fake_response()
    ...

Pytest sam utworzy nowy obiekt dla każdego testu, więc testy nie wpływają na siebie nawzajem.

Jak zasymulować błąd 429 i sprawdzić retry?

Atrapa pozwala wywołać błąd na żądanie. Służy do tego side_effect: gdy podasz listę, kolejne wywołania zwracają kolejne elementy, a elementy będące wyjątkami są rzucane.

from openai import RateLimitError


def rate_limit_error():
    response = MagicMock(status_code=429, headers={})
    return RateLimitError("limit", response=response, body=None)


def test_ask_model_retries_after_429(client):
    client.chat.completions.create.side_effect = [rate_limit_error(), fake_response("OK")]
    sleep = MagicMock()

    assert ask_model(client, "Hej", UsageReport(), sleep=sleep) == "OK"
    assert client.chat.completions.create.call_count == 2
    sleep.assert_called_once_with(1)

Pierwsze wywołanie rzuca RateLimitError, drugie zwraca poprawną odpowiedź. Test sprawdza, że kod ponowił próbę raz i odczekał sekundę, ale prawdziwie nie czeka, bo sleep to atrapa. Drugi test powinien sprawdzić przypadek, gdy błąd się powtarza i kod się poddaje:

import pytest


def test_ask_model_gives_up_after_retries(client):
    client.chat.completions.create.side_effect = rate_limit_error()

    with pytest.raises(RateLimitError):
        ask_model(client, "Hej", UsageReport(), retries=2, sleep=MagicMock())
    assert client.chat.completions.create.call_count == 3

Przy retries=2 mamy trzy próby: pierwszą i dwa ponowienia, dlatego oczekujemy trzech wywołań. Test błędów jest zwykle ważniejszy niż test sukcesu, bo właśnie w ścieżce błędu kod najczęściej zawodzi.

Gdzie użyć patch? Najczęstszy błąd z mockowaniem

Czasem klient nie jest przekazywany jako argument, tylko tworzony wewnątrz modułu. Wtedy używasz patch. Najczęstszy błąd: patchujesz nazwę w miejscu, gdzie została zdefiniowana, a nie tam, gdzie jest używana.

Załóżmy, że plik analyzer.py ma funkcję make_client(), a main() z niej korzysta:

from unittest.mock import patch


def test_main_uses_patched_client(tmp_path, capsys):
    path = tmp_path / "doc.txt"
    path.write_text("Krótki tekst.", encoding="utf-8")
    fake_client = MagicMock()
    fake_client.chat.completions.create.return_value = fake_response("Wynik")

    with patch("analyzer.make_client", return_value=fake_client):
        analyzer.main(str(path))

    assert "Wynik" in capsys.readouterr().out

Zasada jest prosta: patchuj tam, gdzie nazwa jest wyszukiwana. Jeśli w analyzer.py masz from openai import OpenAI, to kod odwołuje się do nazwy analyzer.OpenAI, więc patchujesz "analyzer.OpenAI", a nie "openai.OpenAI". Patch na oryginalnym module nie zadziała, bo analyzer ma już własną kopię nazwy.

Podobny pomysł (podmiana prawdziwych zapytań na sztuczne dane) pojawia się też w JavaScripcie, np. przy limitach API w nauce programowania czy testowaniu kodu asynchronicznego.

Co testować, a czego nie?

Jeśli chcesz kiedyś sprawdzić działanie z prawdziwym API, napisz do tego osobny, rzadko uruchamiany test, oznaczony w pytest (np. markerem), żeby nie startował przy każdej zmianie.

Co dalej?

Znasz już atrapy, side_effect i zasadę dla patch. W ostatniej części wykorzystamy to w kompletnym projekcie: analizatorze dokumentów z raportem zużycia tokenów. Jeśli uczysz się Pythona od podstaw i chcesz robić to z mentorem, zobacz ścieżkę Python.