Add SQLAlchemy multiple binds test

Author: mongkok

tests/panels/sqlalchemy/conftest.py

@@ -3,28 +3,34 @@
 import pytest
 from fastapi import Depends, FastAPI, Request
 from fastapi.responses import HTMLResponse
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 from ...testclient import TestClient
 from .crud import create_user, get_user
-from .database import Base, SessionLocal, engine
+from .database import Base, engine
 
 Base.metadata.create_all(bind=engine)
 
 
-def get_db() -> t.Generator:
-    db = SessionLocal()
-    try:
-        yield db
-    finally:
-        db.close()
+@pytest.fixture
+def get_db(session_options: dict[str, t.Any]) -> t.Callable:
+    SessionLocal = sessionmaker(**session_options or {"bind": engine})
+
+    def f() -> t.Generator:
+        db = SessionLocal()
+        try:
+            yield db
+        finally:
+            db.close()
+
+    return f
 
 
 @pytest.fixture
-def client(app: FastAPI, get_index: t.Callable) -> TestClient:
+def client(app: FastAPI, get_index: t.Callable, get_db: t.Callable) -> TestClient:
     @app.get("/sql", response_class=HTMLResponse)
     async def get_sql(request: Request, db: Session = Depends(get_db)) -> HTMLResponse:
-        user = create_user(db=db, username="test")
+        user = create_user(db=db, username=str(id(get_db)))
         get_user(db=db, user_id=user.id)
         get_user(db=db, user_id=user.id)
         return get_index(request)

tests/panels/sqlalchemy/database.py

@@ -1,11 +1,10 @@
 from sqlalchemy import create_engine
-from sqlalchemy.orm import declarative_base, sessionmaker
+from sqlalchemy.orm import declarative_base
 from sqlalchemy.pool import StaticPool
 
 engine = create_engine(
     "sqlite://",
     connect_args={"check_same_thread": False},
     poolclass=StaticPool,
 )
-SessionLocal = sessionmaker(bind=engine)
 Base = declarative_base()

tests/panels/sqlalchemy/test_sqlalchemy.py

@@ -1,9 +1,19 @@
+import typing as t
+
+import pytest
+from sqlalchemy.orm import declarative_base
+
 from ...mark import override_panels
 from ...testclient import TestClient
+from .database import Base, engine
 
 
 @override_panels(["debug_toolbar.panels.sqlalchemy.SQLAlchemyPanel"])
-def test_sqlalchemy(client: TestClient) -> None:
+@pytest.mark.parametrize(
+    "session_options",
+    (None, {"binds": {Base: engine, declarative_base(): engine}}),
+)
+def test_sqlalchemy(client: TestClient, session_options: dict[str, t.Any]) -> None:
     store_id = client.get_store_id("/sql")
     stats = client.get_stats(store_id, "SQLAlchemyPanel")
     queries = stats["queries"]