feat(models): add more tests on service

.gitignore

@@ -98,3 +98,4 @@ backend/core/examples/chatbot/.chainlit/translations/en-US.json
 
 # Tox
 .tox
+Pipfile

Makefile

@@ -1,15 +1,30 @@
+.DEFAULT_TARGET=help
 
+## help: Display list of commands
+.PHONY: help
+help:
+	@echo "Available commands:"
+	@sed -n 's|^##||p' $(MAKEFILE_LIST) | column -t -s ':' | sed -e 's|^| |'
+
+
+## dev: Start development environment
+.PHONY: dev
 dev:
 	DOCKER_BUILDKIT=1 docker compose -f docker-compose.dev.yml up --build
 
+## prod: Build and start production environment
+.PHONY: prod
 prod:
 	docker compose build backend-core
 	docker compose -f docker-compose.yml up --build
 
-
+## front: Build and start frontend
+.PHONY: front
 front:
 	cd frontend  && yarn build && yarn start
 
+## test: Run tests
+.PHONY: test
 test:
 	# Ensure dependencies are installed with dev and test extras
 	# poetry install --with dev,test && brew install tesseract pandoc libmagic

backend/api/quivr_api/modules/models/tests/conftest.py

@@ -0,0 +1,88 @@
+import asyncio
+import os
+from typing import Tuple
+
+import pytest
+import pytest_asyncio
+import sqlalchemy
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlmodel import select
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from quivr_api.modules.models.entity.model import Model
+from quivr_api.modules.user.entity.user_identity import User
+
+pg_database_base_url = "postgres:postgres@localhost:54322/postgres"
+
+TestData = Tuple[Model, Model, User]
+
+
+@pytest.fixture(scope="session")
+def event_loop(request: pytest.FixtureRequest):
+    loop = asyncio.get_event_loop_policy().new_event_loop()
+    yield loop
+    loop.close()
+
+
+@pytest_asyncio.fixture(scope="session")
+async def async_engine():
+    engine = create_async_engine(
+        "postgresql+asyncpg://" + pg_database_base_url,
+        echo=True if os.getenv("ORM_DEBUG") else False,
+        future=True,
+        pool_pre_ping=True,
+        pool_size=10,
+        pool_recycle=0.1,
+    )
+    yield engine
+
+
+@pytest_asyncio.fixture()
+async def session(async_engine):
+    async with async_engine.connect() as conn:
+        await conn.begin()
+        await conn.begin_nested()
+        async_session = AsyncSession(conn, expire_on_commit=False)
+
+        @sqlalchemy.event.listens_for(
+            async_session.sync_session, "after_transaction_end"
+        )
+        def end_savepoint(session, transaction):
+            if conn.closed:
+                return
+            if not conn.in_nested_transaction():
+                conn.sync_connection.begin_nested()
+
+        yield async_session
+
+
+@pytest_asyncio.fixture()
+async def test_data(
+    session: AsyncSession,
+) -> TestData:
+    # User data
+    user_1 = (
+        await session.exec(select(User).where(User.email == "[email protected]"))
+    ).one()
+
+    model_1 = Model(
+        name="this-is-a-fake-model", price=1, max_input=4000, max_output=2000
+    )
+    model_2 = Model(
+        name="this-is-another-fake-model", price=5, max_input=8000, max_output=4000
+    )
+
+    session.add(model_1)
+    session.add(model_2)
+
+    await session.refresh(user_1)
+    await session.commit()
+    return model_1, model_2, user_1
+
+
+@pytest_asyncio.fixture()
+async def sample_models():
+    return [
+        Model(name="gpt-3.5-turbo", price=1, max_input=4000, max_output=2000),
+        Model(name="gpt-4", price=5, max_input=8000, max_output=4000),
+    ]

backend/api/quivr_api/modules/models/tests/test_models.py

@@ -1,14 +1,6 @@
 import pytest
-import pytest_asyncio
-from quivr_api.modules.models.entity.model import Model
-
 
-@pytest_asyncio.fixture()
-async def sample_models():
-    return [
-        Model(name="gpt-3.5-turbo", price=1, max_input=4000, max_output=2000),
-        Model(name="gpt-4", price=5, max_input=8000, max_output=4000),
-    ]
+from quivr_api.modules.models.entity.model import Model
 
 
 @pytest.mark.asyncio
@@ -21,8 +13,8 @@ async def test_model_creation():
 
 
 @pytest.mark.asyncio
-async def test_model_attributes(sample_models):
-    model = sample_models[0]
+async def test_model_attributes(test_data):
+    model = test_data[0]
     assert hasattr(model, "name")
     assert hasattr(model, "price")
     assert hasattr(model, "max_input")
@@ -66,5 +58,8 @@ async def test_model_dict_representation():
         "price": 2,
         "max_input": 3000,
         "max_output": 1500,
+        "description": "",
+        "image_url": "",
+        "display_name": "",
     }
     assert model.dict() == expected_dict

backend/api/quivr_api/modules/models/tests/test_models_service.py

@@ -0,0 +1,32 @@
+import pytest
+
+from quivr_api.modules.models.repository.model import ModelRepository
+from quivr_api.modules.models.service.model_service import ModelService
+
+
+@pytest.mark.asyncio
+async def test_service_get_chat_models(session):
+    repo = ModelRepository(session)
+    service = ModelService(repo)
+    models = await service.get_models()
+    assert len(models) >= 1
+
+
+@pytest.mark.asyncio
+async def test_service_get_non_existing_chat_model(session):
+    repo = ModelRepository(session)
+    service = ModelService(repo)
+    model = await service.get_model("gpt-3.5-turbo")
+    assert model is None
+
+
+@pytest.mark.asyncio
+async def test_service_get_existing_chat_model(session):
+    repo = ModelRepository(session)
+    service = ModelService(repo)
+    models = await service.get_models()
+    assert len(models) >= 1
+    model = models[0]
+    model_get = await service.get_model(model.name)
+    assert model_get is not None
+    assert model_get == model