From 0686ed1f967166990a940063571582bc861d720d Mon Sep 17 00:00:00 2001
From: John Zaras <48223149+sinnec@users.noreply.github.com>
Date: Thu, 26 Jan 2023 12:49:57 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=A8=E2=80=8D=F0=9F=92=BB=20Drop=20DB?=
 =?UTF-8?q?=20after=20tests=20complete=20(#117)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 whitebox/tests/v1/conftest.py | 42 ++++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 20 deletions(-)

diff --git a/whitebox/tests/v1/conftest.py b/whitebox/tests/v1/conftest.py
index f017683..4e01773 100644
--- a/whitebox/tests/v1/conftest.py
+++ b/whitebox/tests/v1/conftest.py
@@ -1,19 +1,14 @@
-import databases
-import sqlalchemy
 from fastapi.testclient import TestClient
 from pytest import fixture
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import close_all_sessions
 from whitebox import crud
-
 from whitebox.core.settings import get_settings
 from whitebox.entities.Base import Base
 from whitebox.main import app
 from whitebox.sdk.whitebox import Whitebox
 from whitebox.tests.utils.maps import v1_test_order_map
-from whitebox.entities.Base import Base
 from whitebox.utils.passwords import decrypt_api_key
-from whitebox.core.db import SessionLocal
-
+from whitebox.core.db import SessionLocal, engine
 
 settings = get_settings()
 
@@ -28,19 +23,6 @@ def client():
         yield client
 
 
-class TestsState:
-    user: dict = {}
-    model_binary: dict = {}
-    model_multi: dict = {}
-    model_multi_2: dict = {}
-    model_multi_3: dict = {}
-    inference_row_multi: dict = {}
-    inference_row_binary: dict = {}
-
-
-state = TestsState()
-
-
 @fixture(scope="session")
 def api_key():
     db = SessionLocal()
@@ -55,6 +37,26 @@ def api_key():
     yield api_key
 
 
+@fixture(scope="session", autouse=True)
+def drop_db():
+    yield
+    close_all_sessions()
+    Base.metadata.drop_all(engine)
+
+
+class TestsState:
+    user: dict = {}
+    model_binary: dict = {}
+    model_multi: dict = {}
+    model_multi_2: dict = {}
+    model_multi_3: dict = {}
+    inference_row_multi: dict = {}
+    inference_row_binary: dict = {}
+
+
+state = TestsState()
+
+
 class TestsSDKState:
     wb: Whitebox