oterm/tests/conftest.py
2023-12-18 14:18:12 +01:00

26 lines
618 B
Python

from base64 import b64encode
from io import BytesIO
import pytest
import pytest_asyncio
from PIL import Image
from oterm.ollama import OllamaAPI, OllamaError
@pytest_asyncio.fixture(autouse=True)
async def load_test_models():
api = OllamaAPI()
try:
await api.get_model_info("nous-hermes:13b")
except OllamaError:
await api.pull_model("nous-hermes:13b")
yield
@pytest.fixture(scope="session")
def llama_image():
buffered = BytesIO()
image = Image.open("tests/data/lama.jpg")
image.save(buffered, format="JPEG")
return b64encode(buffered.getvalue()).decode("utf-8")