tests

Author: Eugene Shershen

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
     rev: v0.12.4
     hooks:
       - id: ruff
-        args: [--fix]
+        args: [--fix, --unsafe-fixes]
       - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy

tests/test_custom_fastapi_middleware.py

@@ -8,34 +8,47 @@
 import unittest
 from contextvars import ContextVar
 from typing import Dict, Optional, Union
-from unittest.mock import Mock
 
 try:
     from fastapi import FastAPI
-    from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+    from starlette.middleware.base import (
+        BaseHTTPMiddleware,
+        RequestResponseEndpoint,
+    )
     from starlette.requests import Request
     from starlette.types import ASGIApp
+
     FASTAPI_AVAILABLE = True
 except ImportError:
     FASTAPI_AVAILABLE = False
+
     # Create mock classes for when FastAPI is not available
     class FastAPI:
         def __init__(self, *args, **kwargs):
             pass
+
         def add_middleware(self, *args, **kwargs):
             pass
+
         def get(self, *args, **kwargs):
             def decorator(func):
                 return func
+
             return decorator
-    
+
     class BaseHTTPMiddleware:
         def __init__(self, *args, **kwargs):
             pass
 
+
 from sqlalchemy.engine import Engine
 from sqlalchemy.engine.url import URL
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+
 
 # Mock settings for testing
 class MockSettings:
@@ -45,28 +58,40 @@ class MockSettings:
     DB_POOL_SIZE_MAX = 10
     SERVICE_NAME = "test_service"
 
+
 settings = MockSettings()
 
+
 # Mock SSLMode for testing
 class SSLMode:
     disable = "disable"
 
+
 # Custom exceptions for the middleware
 class MissingSessionError(Exception):
     """Raised when no session is available in the current context."""
+
     pass
 
+
 class SessionNotInitialisedError(Exception):
     """Raised when the session factory is not initialized."""
+
     pass
 
 
 def create_middleware_and_session_proxy():
     """Create the custom middleware and session proxy as provided in the issue."""
     _Session: Optional[async_sessionmaker] = None
-    _session: ContextVar[Optional[AsyncSession]] = ContextVar("_session", default=None)
-    _multi_sessions_ctx: ContextVar[bool] = ContextVar("_multi_sessions_context", default=False)
-    _commit_on_exit_ctx: ContextVar[bool] = ContextVar("_commit_on_exit_ctx", default=False)
+    _session: ContextVar[Optional[AsyncSession]] = ContextVar(
+        "_session", default=None
+    )
+    _multi_sessions_ctx: ContextVar[bool] = ContextVar(
+        "_multi_sessions_context", default=False
+    )
+    _commit_on_exit_ctx: ContextVar[bool] = ContextVar(
+        "_commit_on_exit_ctx", default=False
+    )
 
     class SQLAlchemyMiddleware(BaseHTTPMiddleware):
         def __init__(
@@ -84,18 +109,25 @@ def __init__(
             session_args = session_args or {}
 
             if not custom_engine and not db_url:
-                raise ValueError("You need to pass a db_url or a custom_engine parameter.")
+                raise ValueError(
+                    "You need to pass a db_url or a custom_engine parameter."
+                )
             if not custom_engine:
                 engine = create_async_engine(db_url, **engine_args)
             else:
                 engine = custom_engine
 
             nonlocal _Session
             _Session = async_sessionmaker(
-                engine, class_=AsyncSession, expire_on_commit=False, **session_args
+                engine,
+                class_=AsyncSession,
+                expire_on_commit=False,
+                **session_args,
             )
 
-        async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
+        async def dispatch(
+            self, request: Request, call_next: RequestResponseEndpoint
+        ):
             async with DBSession(commit_on_exit=self.commit_on_exit):
                 return await call_next(request)
 
@@ -124,7 +156,9 @@ async def cleanup():
 
                 task = asyncio.current_task()
                 if task is not None:
-                    task.add_done_callback(lambda t: asyncio.create_task(cleanup()))
+                    task.add_done_callback(
+                        lambda t: asyncio.create_task(cleanup())
+                    )
                 return session
             else:
                 session = _session.get()
@@ -151,7 +185,9 @@ async def __aenter__(self):
 
             if self.multi_sessions:
                 self.multi_sessions_token = _multi_sessions_ctx.set(True)
-                self.commit_on_exit_token = _commit_on_exit_ctx.set(self.commit_on_exit)
+                self.commit_on_exit_token = _commit_on_exit_ctx.set(
+                    self.commit_on_exit
+                )
             else:
                 self.token = _session.set(_Session(**self.session_args))
             return type(self)
@@ -185,10 +221,10 @@ def setUp(self):
     def test_middleware_creation_with_db_url(self):
         """Test creating middleware with database URL"""
         SQLAlchemyMiddleware, DBSession = create_middleware_and_session_proxy()
-        
+
         # Create a mock FastAPI app
         app = FastAPI(title="Test App")
-        
+
         # Test middleware creation with basic settings
         try:
             app.add_middleware(
@@ -201,19 +237,24 @@ def test_middleware_creation_with_db_url(self):
                     "max_overflow": settings.DB_POOL_SIZE_MAX,
                     "pool_recycle": 900,
                     "connect_args": {
-                        "server_settings": {"jit": "off", "application_name": settings.SERVICE_NAME},
+                        "server_settings": {
+                            "jit": "off",
+                            "application_name": settings.SERVICE_NAME,
+                        },
                         "ssl": SSLMode.disable,
                     },
                 },
             )
-            self.assertTrue(True, "Custom FastAPI middleware created successfully")
+            self.assertTrue(
+                True, "Custom FastAPI middleware created successfully"
+            )
         except Exception as e:
             self.fail(f"Failed to create custom FastAPI middleware: {e}")
 
     def test_middleware_creation_with_custom_engine(self):
         """Test creating middleware with custom engine"""
         SQLAlchemyMiddleware, DBSession = create_middleware_and_session_proxy()
-        
+
         # Create a custom engine
         engine = create_async_engine(
             settings.DB_URL,
@@ -223,46 +264,57 @@ def test_middleware_creation_with_custom_engine(self):
             max_overflow=settings.DB_POOL_SIZE_MAX,
             pool_recycle=900,
             connect_args={
-                "server_settings": {"jit": "off", "application_name": settings.SERVICE_NAME},
+                "server_settings": {
+                    "jit": "off",
+                    "application_name": settings.SERVICE_NAME,
+                },
                 "ssl": SSLMode.disable,
             },
         )
-        
+
         # Create a mock FastAPI app
         app = FastAPI(title="Test App with Custom Engine")
-        
+
         try:
             app.add_middleware(
                 SQLAlchemyMiddleware,
                 custom_engine=engine,
                 commit_on_exit=True,
             )
-            self.assertTrue(True, "Custom FastAPI middleware with custom engine created successfully")
+            self.assertTrue(
+                True,
+                "Custom FastAPI middleware with custom engine created successfully",
+            )
         except Exception as e:
-            self.fail(f"Failed to create custom FastAPI middleware with custom engine: {e}")
+            self.fail(
+                f"Failed to create custom FastAPI middleware with custom engine: {e}"
+            )
 
     def test_middleware_validation_errors(self):
         """Test middleware validation for required parameters"""
         SQLAlchemyMiddleware, DBSession = create_middleware_and_session_proxy()
-        
+
         app = FastAPI(title="Test App")
-        
+
         # Test missing both db_url and custom_engine by calling the constructor directly
         with self.assertRaises(ValueError) as context:
             # This should raise ValueError when called directly
-            middleware = SQLAlchemyMiddleware(app)
-        
-        self.assertIn("You need to pass a db_url or a custom_engine parameter", str(context.exception))
+            SQLAlchemyMiddleware(app)
+
+        self.assertIn(
+            "You need to pass a db_url or a custom_engine parameter",
+            str(context.exception),
+        )
 
     def test_db_session_context_manager(self):
         """Test DBSession context manager functionality"""
         SQLAlchemyMiddleware, DBSession = create_middleware_and_session_proxy()
-        
+
         # Test DBSession creation
         db_session = DBSession(commit_on_exit=True)
         self.assertIsInstance(db_session, DBSession)
         self.assertTrue(db_session.commit_on_exit)
-        
+
         # Test multi-sessions mode
         multi_db_session = DBSession(multi_sessions=True, commit_on_exit=False)
         self.assertIsInstance(multi_db_session, DBSession)
@@ -272,26 +324,30 @@ def test_db_session_context_manager(self):
     def test_session_proxy_creation(self):
         """Test that the session proxy is created correctly"""
         SQLAlchemyMiddleware, DBSession = create_middleware_and_session_proxy()
-        
+
         # Verify that we get the expected classes
         if FASTAPI_AVAILABLE:
-            self.assertTrue(issubclass(SQLAlchemyMiddleware, BaseHTTPMiddleware))
-        
+            self.assertTrue(
+                issubclass(SQLAlchemyMiddleware, BaseHTTPMiddleware)
+            )
+
         # Test that DBSession has the expected metaclass behavior
-        self.assertTrue(hasattr(DBSession, '__aenter__'))
-        self.assertTrue(hasattr(DBSession, '__aexit__'))
-        
+        self.assertTrue(hasattr(DBSession, "__aenter__"))
+        self.assertTrue(hasattr(DBSession, "__aexit__"))
+
         # Test that the session property exists on the class (without accessing it)
         # We check the metaclass has the session property
-        self.assertTrue(hasattr(type(DBSession), 'session'))
-        self.assertTrue(isinstance(getattr(type(DBSession), 'session'), property))
+        self.assertTrue(hasattr(type(DBSession), "session"))
+        self.assertTrue(
+            isinstance(getattr(type(DBSession), "session"), property)
+        )
 
     def test_middleware_with_session_args(self):
         """Test middleware creation with session arguments"""
         SQLAlchemyMiddleware, DBSession = create_middleware_and_session_proxy()
-        
+
         app = FastAPI(title="Test App with Session Args")
-        
+
         try:
             app.add_middleware(
                 SQLAlchemyMiddleware,
@@ -306,17 +362,22 @@ def test_middleware_with_session_args(self):
                 },
                 commit_on_exit=True,
             )
-            self.assertTrue(True, "Custom FastAPI middleware with session args created successfully")
+            self.assertTrue(
+                True,
+                "Custom FastAPI middleware with session args created successfully",
+            )
         except Exception as e:
-            self.fail(f"Failed to create custom FastAPI middleware with session args: {e}")
+            self.fail(
+                f"Failed to create custom FastAPI middleware with session args: {e}"
+            )
 
     @unittest.skipUnless(FASTAPI_AVAILABLE, "FastAPI not available")
     def test_full_integration_mock(self):
         """Test full integration with mocked database operations"""
         SQLAlchemyMiddleware, DBSession = create_middleware_and_session_proxy()
-        
+
         app = FastAPI(title="Full Integration Test")
-        
+
         # Add the middleware
         app.add_middleware(
             SQLAlchemyMiddleware,
@@ -328,27 +389,33 @@ def test_full_integration_mock(self):
                 "max_overflow": settings.DB_POOL_SIZE_MAX,
                 "pool_recycle": 900,
                 "connect_args": {
-                    "server_settings": {"jit": "off", "application_name": settings.SERVICE_NAME},
+                    "server_settings": {
+                        "jit": "off",
+                        "application_name": settings.SERVICE_NAME,
+                    },
                     "ssl": SSLMode.disable,
                 },
             },
             commit_on_exit=True,
         )
-        
+
         # Add a test route
         @app.get("/test")
         async def test_route():
             return {"message": "Test successful"}
-        
+
         # Test that the app and middleware were set up successfully
         # We don't need to actually make HTTP requests, just verify setup
         self.assertIsNotNone(app)
-        self.assertTrue(True, "Full integration test completed - middleware setup successful")
+        self.assertTrue(
+            True,
+            "Full integration test completed - middleware setup successful",
+        )
 
 
 if __name__ == "__main__":
     if FASTAPI_AVAILABLE:
         unittest.main()
     else:
         print("FastAPI not available. Skipping tests.")
-        print("To run these tests, install FastAPI: pip install fastapi")
+        print("To run these tests, install FastAPI: pip install fastapi")