chore: cleanup and split up tests

Author: cofin

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
       - id: unasyncd
         additional_dependencies: ["ruff"]
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.14.14"
+    rev: "v0.15.0"
     hooks:
       # Run the linter.
       - id: ruff

advanced_alchemy/alembic/utils.py

@@ -6,6 +6,7 @@
 from typing_extensions import TypeIs
 
 from advanced_alchemy.exceptions import MissingDependencyError
+from advanced_alchemy.utils.sync_tools import async_
 
 if TYPE_CHECKING:
     from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
@@ -133,8 +134,8 @@ async def _dump_table_async(session: "AbstractAsyncContextManager[AsyncSession]"
                 )
                 json_path.write_text(encode_json([row.to_dict() for row in await repo(session=_session).list()]))
 
-    dump_dir.mkdir(exist_ok=True)
+    await async_(dump_dir.mkdir)(exist_ok=True)
 
     if _is_sync(session):
-        return _dump_table_sync(session)
+        return await async_(_dump_table_sync)(session)
     return await _dump_table_async(session)

advanced_alchemy/repository/memory/base.py

@@ -313,7 +313,7 @@ class Child(Base):
                 continue
             remote_mapper = mappers[next(iter(column.foreign_keys))._table_key()]  # noqa: SLF001 # pyright: ignore[reportPrivateUsage]
             try:
-                obj = self.store(remote_mapper.class_).get(new_attrs.get(column.key, None))
+                obj = self.store(remote_mapper.class_).get(new_attrs.get(column.key))
             except KeyError:
                 continue
 

advanced_alchemy/types/file_object/backends/obstore.py

@@ -142,7 +142,7 @@ def save_object(
             "max_concurrency": max_concurrency,
         }
         if not isinstance(self.fs, LocalStore):
-            put_params["attributes"] = attributes if attributes else None
+            put_params["attributes"] = attributes or None
 
         _ = self.fs.put(file_object.path, data, **put_params)
         info = self.fs.head(file_object.path)
@@ -198,7 +198,7 @@ async def save_object_async(
             "max_concurrency": max_concurrency,
         }
         if not isinstance(self.fs, LocalStore):
-            put_params["attributes"] = attributes if attributes else None
+            put_params["attributes"] = attributes or None
 
         _ = await self.fs.put_async(file_object.path, data, **put_params)
         info = await self.fs.head_async(file_object.path)

tests/integration/test_alembic_commands.py

@@ -223,23 +223,27 @@ def tmp_project_dir(monkeypatch: MonkeyPatch, tmp_path: Path) -> Generator[Path,
 
 
 async def test_alembic_init(alembic_commands: commands.AlembicCommands, tmp_project_dir: Path) -> None:
+    from advanced_alchemy.utils.sync_tools import async_
+
     alembic_commands.init(directory=f"{tmp_project_dir}/migrations/")
     expected_dirs = [f"{tmp_project_dir}/migrations/", f"{tmp_project_dir}/migrations/versions"]
     expected_files = [f"{tmp_project_dir}/migrations/env.py", f"{tmp_project_dir}/migrations/script.py.mako"]
     for dir in expected_dirs:
-        assert Path(dir).is_dir()
+        assert await async_(Path(dir).is_dir)()
     for file in expected_files:
-        assert Path(file).is_file()
+        assert await async_(Path(file).is_file)()
 
 
 async def test_alembic_init_already(alembic_commands: commands.AlembicCommands, tmp_project_dir: Path) -> None:
+    from advanced_alchemy.utils.sync_tools import async_
+
     alembic_commands.init(directory=f"{tmp_project_dir}/migrations/")
     expected_dirs = [f"{tmp_project_dir}/migrations/", f"{tmp_project_dir}/migrations/versions"]
     expected_files = [f"{tmp_project_dir}/migrations/env.py", f"{tmp_project_dir}/migrations/script.py.mako"]
     for dir in expected_dirs:
-        assert Path(dir).is_dir()
+        assert await async_(Path(dir).is_dir)()
     for file in expected_files:
-        assert Path(file).is_file()
+        assert await async_(Path(file).is_file)()
     with pytest.raises(CommandError):
         alembic_commands.init(directory=f"{tmp_project_dir}/migrations/")
 

tests/integration/test_file_object.py

@@ -959,7 +959,9 @@ async def test_file_object_save_with_different_data_types(storage_registry: Stor
     assert await obj2.get_content_async() == test_content
     assert obj2.get_content() == test_content
     # Cleanup
-    temp_path.unlink()
+    from advanced_alchemy.utils.sync_tools import async_
+
+    await async_(temp_path.unlink)()
 
 
 @pytest.mark.xdist_group("file_object")

tests/integration/test_operations.py

@@ -120,7 +120,7 @@ async def test_table_async(
     yield cached_test_table
 
     # Fast data-only cleanup between tests
-    if dialect_name not in {"mock"} and "spanner" not in dialect_name:
+    if dialect_name != "mock" and "spanner" not in dialect_name:
         async with async_engine.begin() as conn:
             await conn.execute(cached_test_table.delete())
             await conn.commit()
@@ -255,7 +255,7 @@ def store_model_sync(
     yield cached_store_model
 
     # Fast data-only cleanup between tests
-    if dialect_name not in {"mock"} and "spanner" not in dialect_name:
+    if dialect_name != "mock" and "spanner" not in dialect_name:
         clean_tables(engine, cached_store_model.metadata)
 
     # Drop table at session end (handled by teardown)
@@ -296,7 +296,7 @@ async def store_model_async(
     yield cached_store_model
 
     # Fast data-only cleanup between tests
-    if dialect_name not in {"mock"} and "spanner" not in dialect_name:
+    if dialect_name != "mock" and "spanner" not in dialect_name:
         await async_clean_tables(async_engine, cached_store_model.metadata)
 
     # Drop table at session end (handled by teardown)

tests/unit/test_cache/test_cache_listeners.py

@@ -0,0 +1,214 @@
+"""Unit tests for cache-related listeners in advanced_alchemy._listeners."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from sqlalchemy.orm import Session
+
+from advanced_alchemy._listeners import (
+    AsyncCacheListener,
+    BaseCacheListener,
+    CacheInvalidationListener,
+    CacheInvalidationTracker,
+    SyncCacheListener,
+    get_cache_tracker,
+    setup_cache_listeners,
+)
+
+# --- CacheInvalidationTracker Tests ---
+
+
+def test_cache_invalidation_tracker_add_invalidation() -> None:
+    mock_manager = MagicMock()
+    tracker = CacheInvalidationTracker(mock_manager)
+    tracker.add_invalidation("User", 1, "group1")
+
+    assert ("User", 1, "group1") in tracker._pending_invalidations
+    assert "User" in tracker._pending_model_bumps
+
+
+def test_cache_invalidation_tracker_commit() -> None:
+    mock_manager = MagicMock()
+    tracker = CacheInvalidationTracker(mock_manager)
+    tracker.add_invalidation("User", 1, "group1")
+
+    tracker.commit()
+
+    mock_manager.bump_model_version_sync.assert_called_with("User")
+    mock_manager.invalidate_entity_sync.assert_called_with("User", 1, "group1")
+    assert not tracker._pending_invalidations
+    assert not tracker._pending_model_bumps
+
+
+def test_cache_invalidation_tracker_rollback() -> None:
+    mock_manager = MagicMock()
+    tracker = CacheInvalidationTracker(mock_manager)
+    tracker.add_invalidation("User", 1, "group1")
+
+    tracker.rollback()
+
+    mock_manager.bump_model_version_sync.assert_not_called()
+    assert not tracker._pending_invalidations
+    assert not tracker._pending_model_bumps
+
+
+@pytest.mark.asyncio
+async def test_cache_invalidation_tracker_commit_async() -> None:
+    mock_manager = AsyncMock()
+    tracker = CacheInvalidationTracker(mock_manager)
+    tracker.add_invalidation("User", 1, "group1")
+
+    await tracker.commit_async()
+
+    mock_manager.bump_model_version_async.assert_called_with("User")
+    mock_manager.invalidate_entity_async.assert_called_with("User", 1, "group1")
+    assert not tracker._pending_invalidations
+    assert not tracker._pending_model_bumps
+
+
+# --- get_cache_tracker Tests ---
+
+
+def test_get_cache_tracker_existing() -> None:
+    session = MagicMock(spec=Session)
+    tracker = MagicMock()
+    session.info = {"_aa_cache_tracker": tracker}
+
+    result = get_cache_tracker(session)
+    assert result is tracker
+
+
+def test_get_cache_tracker_create() -> None:
+    session = MagicMock(spec=Session)
+    session.info = {}
+    mock_manager = MagicMock()
+
+    result = get_cache_tracker(session, cache_manager=mock_manager, create=True)
+    assert isinstance(result, CacheInvalidationTracker)
+    assert session.info["_aa_cache_tracker"] is result
+
+
+def test_get_cache_tracker_no_create() -> None:
+    session = MagicMock(spec=Session)
+    session.info = {}
+
+    result = get_cache_tracker(session, create=False)
+    assert result is None
+
+
+# --- SyncCacheListener Tests ---
+
+
+def test_sync_cache_listener_after_commit() -> None:
+    session = MagicMock(spec=Session)
+    tracker = MagicMock()
+    session.info = {"_aa_cache_tracker": tracker, "enable_cache_listener": True}
+
+    SyncCacheListener.after_commit(session)
+
+    tracker.commit.assert_called_once()
+    assert "_aa_cache_tracker" not in session.info
+
+
+def test_sync_cache_listener_after_rollback() -> None:
+    session = MagicMock(spec=Session)
+    tracker = MagicMock()
+    session.info = {"_aa_cache_tracker": tracker, "enable_cache_listener": True}
+
+    SyncCacheListener.after_rollback(session)
+
+    tracker.rollback.assert_called_once()
+    assert "_aa_cache_tracker" not in session.info
+
+
+def test_sync_cache_listener_disabled() -> None:
+    session = MagicMock(spec=Session)
+    tracker = MagicMock()
+    session.info = {"_aa_cache_tracker": tracker, "enable_cache_listener": False}
+
+    SyncCacheListener.after_commit(session)
+
+    tracker.commit.assert_not_called()
+
+
+# --- AsyncCacheListener Tests ---
+
+
+@pytest.mark.asyncio
+async def test_async_cache_listener_after_commit() -> None:
+    session = MagicMock(spec=Session)
+    tracker = MagicMock()
+    # Mock commit_async to verify it's scheduled
+    tracker.commit_async = AsyncMock()
+    session.info = {"_aa_cache_tracker": tracker, "enable_cache_listener": True}
+
+    AsyncCacheListener.after_commit(session)
+
+    # Since it creates a task, we verify it popped the tracker
+    assert "_aa_cache_tracker" not in session.info
+
+
+def test_async_cache_listener_after_rollback() -> None:
+    session = MagicMock(spec=Session)
+    tracker = MagicMock()
+    session.info = {"_aa_cache_tracker": tracker, "enable_cache_listener": True}
+
+    AsyncCacheListener.after_rollback(session)
+
+    tracker.rollback.assert_called_once()
+    assert "_aa_cache_tracker" not in session.info
+
+
+# --- CacheInvalidationListener Tests ---
+
+
+def test_cache_invalidation_listener_after_commit_sync_context() -> None:
+    session = MagicMock(spec=Session)
+    tracker = MagicMock()
+    session.info = {"_aa_cache_tracker": tracker, "enable_cache_listener": True}
+
+    try:
+        CacheInvalidationListener.after_commit(session)
+    except RuntimeError:
+        # If no loop, it calls tracker.commit()
+        pass
+
+
+# --- BaseCacheListener Tests ---
+
+
+class TestBaseCacheListener(BaseCacheListener):
+    pass
+
+
+def test_base_cache_listener_is_listener_enabled() -> None:
+    session = MagicMock(spec=Session)
+    session.info = {}
+    session.bind = None
+    session.execution_options = None
+    assert TestBaseCacheListener._is_listener_enabled(session) is True
+
+    session.info = {"enable_cache_listener": False}
+    assert TestBaseCacheListener._is_listener_enabled(session) is False
+
+
+def test_base_cache_listener_execution_options() -> None:
+    session = MagicMock(spec=Session)
+    session.info = {}
+    session.bind = None
+
+    # Test execution_options via session
+    session.execution_options = {"enable_cache_listener": False}
+    assert TestBaseCacheListener._is_listener_enabled(session) is False
+
+    session.execution_options = {"enable_cache_listener": True}
+    assert TestBaseCacheListener._is_listener_enabled(session) is True
+
+
+# --- Setup Tests ---
+
+
+def test_setup_cache_listeners() -> None:
+    with patch("sqlalchemy.event.listen") as mock_listen:
+        setup_cache_listeners()
+        assert mock_listen.called