From eb182b71f47397bfc9b40bd89d1fcd2c3ee45bd7 Mon Sep 17 00:00:00 2001
From: stijndcl <declercq.stijn@outlook.com>
Date: Fri, 17 Jun 2022 01:02:36 +0200
Subject: [PATCH] Create connection fixture

---
 pyproject.toml    |  2 +-
 tests/conftest.py | 26 ++++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index aada43c..bf0e7fd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,4 +24,4 @@ env = [
     "DB_USERNAME = postgres",
     "DB_HOST = localhost",
     "DISC_TOKEN = token"
-]
\ No newline at end of file
+]
diff --git a/tests/conftest.py b/tests/conftest.py
index 40d3189..37e249c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,6 +1,12 @@
+import os
+from typing import AsyncGenerator
+
 import pytest
 
 from alembic import command, config
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from database.engine import engine
 
 
 @pytest.fixture(scope="session")
@@ -9,7 +15,27 @@ def tables():
     Starts from an empty database and runs through all the migrations to check those as well
     while we're at it
     """
+    print("CWD: ", os.getcwd())
     alembic_config = config.Config("alembic.ini")
     command.upgrade(alembic_config, "head")
     yield
     command.downgrade(alembic_config, "base")
+
+
+@pytest.fixture
+async def database_session(tables) -> AsyncGenerator[AsyncSession, None]:
+    """Fixture to create a session for every test
+    Rollbacks the transaction afterwards so that the future tests start with a clean database
+    """
+    connection = await engine.connect()
+    transaction = await connection.begin()
+    session = AsyncSession(bind=connection, expire_on_commit=False)
+
+    yield session
+
+    # Clean up session & rollback transactions
+    await session.close()
+    if transaction.is_valid:
+        await transaction.rollback()
+
+    await connection.close()