From 44cf716c83a8123f85965d3ba24248d5020cf15a Mon Sep 17 00:00:00 2001
From: Devansh Singh <devanshamity@gmail.com>
Date: Mon, 1 Jul 2024 20:12:00 +0530
Subject: [PATCH] tests/conftest.py: Add fixtures

Add fixtures to conftest.py for database session and FastAPI
TestClient

Signed-off-by: Devansh Singh <devanshamity@gmail.com>
---
 tests/conftest.py     | 39 ++++++++++++++++++++++++++++++++++++++-
 tests/test_helpers.py |  6 +-----
 tests/test_kill.py    | 18 ++++--------------
 tests/test_suite.py   | 18 +++++-------------
 4 files changed, 48 insertions(+), 33 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 8f9e147..326dd05 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -7,4 +7,41 @@
     - https://docs.pytest.org/en/stable/writing_plugins.html
 """
 
-# import pytest
+import pytest
+
+from fastapi.testclient import TestClient
+from sqlmodel import Session, SQLModel, create_engine
+from sqlmodel.pool import StaticPool
+
+from teuthology_api.main import app
+from teuthology_api.models import get_db
+from teuthology_api.services.helpers import get_token
+
+
+def override_get_token():
+    return {"access_token": "token_123", "token_type": "bearer"}
+
+
+@pytest.fixture(name="session", scope="session")
+def session_fixture():
+    engine = create_engine(
+        "sqlite://",
+        connect_args={"check_same_thread": False},
+        poolclass=StaticPool,
+    )
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        yield session
+
+
+@pytest.fixture(name="client", scope="session")
+def client_fixture(session: Session):
+    def override_get_db():
+        return session
+
+    app.dependency_overrides[get_db] = override_get_db
+    app.dependency_overrides[get_token] = override_get_token
+
+    client = TestClient(app)
+    yield client
+    app.dependency_overrides.clear()
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
index cda9a64..dad59d7 100644
--- a/tests/test_helpers.py
+++ b/tests/test_helpers.py
@@ -1,11 +1,7 @@
-from fastapi.testclient import TestClient
 from fastapi import HTTPException
 import pytest
-from teuthology_api.main import app
 from unittest.mock import patch
-from teuthology_api.services.helpers import Request, get_token, get_username
-
-client = TestClient(app)
+from teuthology_api.services.helpers import get_token, get_username
 
 
 class MockRequest:
diff --git a/tests/test_kill.py b/tests/test_kill.py
index 7a30953..9de696b 100644
--- a/tests/test_kill.py
+++ b/tests/test_kill.py
@@ -1,19 +1,7 @@
 from fastapi.testclient import TestClient
-from teuthology_api.main import app
 from unittest.mock import patch
-from teuthology_api.services.helpers import get_token
-from teuthology_api.services.kill import get_username, get_run_details
 import json
-from teuthology_api.schemas.kill import KillArgs
 
-client = TestClient(app)
-
-
-async def override_get_token():
-    return {"access_token": "token_123", "token_type": "bearer"}
-
-
-app.dependency_overrides[get_token] = override_get_token
 
 mock_kill_args = {
     "--dry-run": False,
@@ -33,7 +21,9 @@ async def override_get_token():
 @patch("subprocess.Popen")
 @patch("teuthology_api.services.kill.get_run_details")
 @patch("teuthology_api.services.kill.get_username")
-def test_kill_run_success(m_get_username, m_get_run_details, m_popen):
+def test_kill_run_success(
+    m_get_username, m_get_run_details, m_popen, client: TestClient
+):
     m_get_username.return_value = "user1"
     m_get_run_details.return_value = {"id": "7451978", "user": "user1"}
     mock_process = m_popen.return_value
@@ -44,7 +34,7 @@ def test_kill_run_success(m_get_username, m_get_run_details, m_popen):
     assert response.json() == {"kill": "success"}
 
 
-def test_kill_run_fail():
+def test_kill_run_fail(client: TestClient):
     response = client.post("/kill", data=json.dumps(mock_kill_args))
     assert response.status_code == 401
     assert response.json() == {"detail": "You need to be logged in"}
diff --git a/tests/test_suite.py b/tests/test_suite.py
index ab2f6f7..3117bc9 100644
--- a/tests/test_suite.py
+++ b/tests/test_suite.py
@@ -1,19 +1,8 @@
 from fastapi.testclient import TestClient
-from teuthology_api.main import app
 from unittest.mock import patch
-from teuthology_api.services.helpers import get_token
-from teuthology_api.services.suite import make_run_name, get_run_details
+from teuthology_api.services.suite import make_run_name
 import json
 
-client = TestClient(app)
-
-
-async def override_get_token():
-    return {"access_token": "token_123", "token_type": "bearer"}
-
-
-app.dependency_overrides[get_token] = override_get_token
-
 mock_suite_args = {
     "--dry-run": False,
     "--non-interactive": False,
@@ -28,11 +17,14 @@ async def override_get_token():
     "--machine-type": "testnode",
 }
 
+
 # suite
 @patch("teuthology_api.services.suite.logs_run")
 @patch("teuthology_api.routes.suite.get_username")
 @patch("teuthology_api.services.suite.get_run_details")
-def test_suite_run_success(m_get_run_details, m_get_username, m_logs_run):
+def test_suite_run_success(
+    m_get_run_details, m_get_username, m_logs_run, client: TestClient
+):
     m_get_username.return_value = "user1"
     m_get_run_details.return_value = {"id": "7451978", "user": "user1"}
     response = client.post("/suite", data=json.dumps(mock_suite_args))