👨‍💻 Drop DB after tests complete (#117)

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