Test db setup helpers

Author: chriscarrollsmith

tests/test_db.py

@@ -1,11 +1,15 @@
-from sqlmodel import Session, select
+import warnings
+from sqlmodel import Session, select, inspect
 from utils.db import (
     get_connection_url,
     assign_permissions_to_role,
     create_default_roles,
     create_permissions,
+    tear_down_db,
+    set_up_db,
 )
 from utils.models import Role, Permission, Organization, RolePermissionLink, ValidPermissions
+from sqlalchemy import create_engine
 
 
 def test_get_connection_url():
@@ -120,3 +124,66 @@ def test_assign_permissions_to_role_duplicate_check(session: Session, test_organ
         )
     ).all()
     assert len(link_count) == 1
+
+
+def test_set_up_db_creates_tables():
+    """Test that set_up_db creates all expected tables without warnings"""
+    # First tear down any existing tables
+    tear_down_db()
+
+    # Run set_up_db with drop=False since we just cleaned up
+    set_up_db(drop=False)
+
+    # Use SQLAlchemy inspect to check tables
+    engine = create_engine(get_connection_url())
+    inspector = inspect(engine)
+    table_names = inspector.get_table_names()
+
+    # Check for core tables
+    expected_tables = {
+        "user",
+        "organization",
+        "role",
+        "permission",
+        "role_permission_link",
+        "password_reset_token"
+    }
+    assert expected_tables.issubset(set(table_names))
+
+    # Verify permissions were created
+    with Session(engine) as session:
+        permissions = session.exec(select(Permission)).all()
+        assert len(permissions) == len(ValidPermissions)
+
+    # Clean up
+    tear_down_db()
+    engine.dispose()
+
+
+def test_set_up_db_drop_flag():
+    """Test that set_up_db's drop flag properly recreates tables"""
+    # Set up db with drop=True
+    engine = create_engine(get_connection_url())
+    set_up_db(drop=True)
+
+    # Create a new session for this test
+    with Session(engine) as session:
+        # Verify valid permissions exist
+        permissions = session.exec(select(Permission)).all()
+        assert len(permissions) == len(ValidPermissions)
+
+        # Create an organization
+        org = Organization(name="Test Organization")
+        session.add(org)
+        session.commit()
+
+        # Set up db with drop=False
+        set_up_db(drop=False)
+
+        # Verify organization exists
+        assert session.exec(select(Organization).where(
+            Organization.name == "Test Organization")).first() is not None
+
+    # Clean up
+    tear_down_db()
+    engine.dispose()