From e5d74d99571a5e3b27cd6cd86ea0aaae758ab02e Mon Sep 17 00:00:00 2001
From: GwynBleidD <gbd.lin@gmail.com>
Date: Thu, 8 Aug 2024 23:34:51 +0200
Subject: [PATCH] Mypy fixes

---
 Makefile            | 4 ++--
 tests/conftest.py   | 2 +-
 tests/db_helpers.py | 9 ++++++++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/Makefile b/Makefile
index e35404e..23e4370 100644
--- a/Makefile
+++ b/Makefile
@@ -34,10 +34,10 @@ migrations/create:
 	alembic revision
 
 dc/build:
-	docker-compose -f docker-compose.local.yml build
+	docker compose -f docker-compose.local.yml build
 
 dc/%:
-	docker-compose -f docker-compose.local.yml run -w /app plugin_store make $*
+	docker compose -f docker-compose.local.yml run -w /app plugin_store make $*
 
 deps/lock:
 	poetry lock --no-update
diff --git a/tests/conftest.py b/tests/conftest.py
index 1d9476a..f212fb9 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -82,7 +82,7 @@ def seed_db_sessionmaker(seed_db_engine: tuple["AsyncEngine", "sessionmaker"]) -
 
 
 @pytest_asyncio.fixture()
-async def seed_db_session(seed_db_engine: tuple["AsyncEngine", "sessionmaker"]) -> "AsyncSession":
+async def seed_db_session(seed_db_engine: tuple["AsyncEngine", "sessionmaker"]) -> "AsyncIterator[AsyncSession]":
     async with prepare_transactioned_db_session(*seed_db_engine) as session:
         yield session
 
diff --git a/tests/db_helpers.py b/tests/db_helpers.py
index 29f121d..86512f2 100644
--- a/tests/db_helpers.py
+++ b/tests/db_helpers.py
@@ -3,6 +3,7 @@
 from datetime import datetime, timedelta, UTC
 from hashlib import sha256
 from os import getenv
+from typing import TYPE_CHECKING
 
 from sqlalchemy import event
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
@@ -11,6 +12,9 @@
 
 from database.models import Artifact, Base, Tag, Version
 
+if TYPE_CHECKING:
+    from typing import AsyncIterator
+
 
 class FakePluginGenerator:
     def __init__(
@@ -168,7 +172,10 @@ async def prepare_test_db(
 
 
 @asynccontextmanager
-async def prepare_transactioned_db_session(engine: "AsyncEngine", db_sessionmaker: "sessionmaker") -> "AsyncSession":
+async def prepare_transactioned_db_session(
+    engine: "AsyncEngine",
+    db_sessionmaker: "sessionmaker",
+) -> "AsyncIterator[AsyncSession]":
     connection = await engine.connect()
     outer_transaction = await connection.begin()
     async_session = db_sessionmaker(bind=connection)