"""Step definitions for UnitOfWork coverage boost. Targets uncovered lines in ``src/cleveragents/infrastructure/database/unit_of_work.py``: - Lines 71-76: fresh in-memory engine creation (MEMORY_ENGINES cache miss) - Lines 90-93: non-SQLite engine creation - Lines 244-248: checkpoints property lazy init - Line 312: refresh() method """ from __future__ import annotations from unittest.mock import MagicMock, patch from behave import given, then, when from behave.runner import Context from cleveragents.infrastructure.database.unit_of_work import ( UnitOfWork, UnitOfWorkContext, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_mock_session() -> MagicMock: """Create a MagicMock that mimics a SQLAlchemy Session.""" session = MagicMock() session.add = MagicMock() session.flush = MagicMock() session.refresh = MagicMock() session.commit = MagicMock() session.rollback = MagicMock() session.close = MagicMock() return session # ========================================================================= # Non-SQLite engine (lines 90-93) # ========================================================================= @given("a UnitOfWork configured with a non-SQLite database URL") def step_uow_non_sqlite_url(context: Context) -> None: """Create a UnitOfWork with a PostgreSQL-style URL (never actually connects).""" context.uow = UnitOfWork( database_url="postgresql://user:pass@localhost:5432/testdb" ) @when("I access the engine property with mocked create_engine") def step_access_engine_property_mocked(context: Context) -> None: """Access the engine property, mocking _ensure_database_initialized and create_engine.""" import sys as _sys from sqlalchemy.engine import Engine mock_engine = MagicMock(spec=Engine) # Ensure psycopg2 can be imported (it may not be installed in test env) _needs_psycopg2_stub = "psycopg2" not in _sys.modules if _needs_psycopg2_stub: _sys.modules["psycopg2"] = MagicMock() try: with ( patch.object(context.uow, "_ensure_database_initialized"), patch( "cleveragents.infrastructure.database.unit_of_work.create_engine", return_value=mock_engine, ) as mock_create, ): engine = context.uow.engine context.result_engine = engine context.mock_create_engine = mock_create finally: if _needs_psycopg2_stub: _sys.modules.pop("psycopg2", None) @then("the engine should be created via the non-SQLite code path") def step_engine_created_non_sqlite(context: Context) -> None: """Verify that create_engine was called (meaning the non-SQLite branch ran).""" assert context.mock_create_engine.called, "create_engine was not called" @then("the engine should not use SQLite-specific connect_args") def step_engine_no_sqlite_args(context: Context) -> None: """Verify the call did not include connect_args or isolation_level.""" call_kwargs = context.mock_create_engine.call_args[1] assert "connect_args" not in call_kwargs, ( f"connect_args should not be present for non-SQLite, got {call_kwargs}" ) assert "isolation_level" not in call_kwargs, ( f"isolation_level should not be present for non-SQLite, got {call_kwargs}" ) # ========================================================================= # Fresh in-memory engine (lines 71-76) # ========================================================================= @given("the in-memory engine cache is cleared") def step_clear_memory_engine_cache(context: Context) -> None: """Clear the MEMORY_ENGINES cache so the cache-miss branch is exercised.""" from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES # Save and clear so the test exercises the cache-miss path context.saved_memory_engines = dict(MEMORY_ENGINES) MEMORY_ENGINES.clear() @given("a UnitOfWork configured with an in-memory SQLite URL") def step_uow_in_memory_sqlite(context: Context) -> None: """Create a UnitOfWork targeting in-memory SQLite.""" context.uow = UnitOfWork(database_url="sqlite:///:memory:") @when("I access the in-memory engine property with mocked migrations") def step_access_in_memory_engine_mocked_migrations(context: Context) -> None: """Access the engine property, mocking only _ensure_database_initialized. The real create_engine runs so that MEMORY_ENGINES is actually populated with a real engine, exercising lines 71-76. """ with patch.object(context.uow, "_ensure_database_initialized"): engine = context.uow.engine context.result_engine = engine @then("a new in-memory engine should be created and cached") def step_engine_created_and_cached(context: Context) -> None: """Verify the engine was stored in MEMORY_ENGINES.""" from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES assert "sqlite:///:memory:" in MEMORY_ENGINES, ( "Expected sqlite:///:memory: in MEMORY_ENGINES cache" ) assert context.result_engine is MEMORY_ENGINES["sqlite:///:memory:"], ( "Engine returned should be the same as the cached one" ) @then("the cached engine should use SERIALIZABLE isolation") def step_cached_engine_serializable(context: Context) -> None: """Verify the created engine is cached and has the correct URL. The SERIALIZABLE isolation level is set at creation time. We verify the engine is present and has the correct URL as a proxy. """ from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES engine = MEMORY_ENGINES.get("sqlite:///:memory:") assert engine is not None, "Engine should be present in cache" assert str(engine.url) == "sqlite:///:memory:", ( f"Expected sqlite:///:memory: URL, got {engine.url}" ) # Restore saved cache state if hasattr(context, "saved_memory_engines"): MEMORY_ENGINES.clear() MEMORY_ENGINES.update(context.saved_memory_engines) # ========================================================================= # Checkpoint repository (lines 244-248) # ========================================================================= @given("a UnitOfWorkContext backed by a mock session") def step_uow_context_mock_session(context: Context) -> None: """Create a UnitOfWorkContext with a mock session.""" context.mock_session = _make_mock_session() context.uow_ctx = UnitOfWorkContext(context.mock_session) @when("I access the checkpoints repository from the context") def step_access_checkpoints_repo(context: Context) -> None: """Access the checkpoints property.""" context.checkpoints_result = context.uow_ctx.checkpoints @then("the checkpoints repository should be a CheckpointRepository instance") def step_checkpoints_is_correct_type(context: Context) -> None: """Verify the result is a CheckpointRepository.""" from cleveragents.infrastructure.database.repositories import ( CheckpointRepository, ) assert isinstance(context.checkpoints_result, CheckpointRepository), ( f"Expected CheckpointRepository, got {type(context.checkpoints_result)}" ) @when("I access the checkpoints repository from the context twice") def step_access_checkpoints_repo_twice(context: Context) -> None: """Access the checkpoints property twice.""" context.checkpoints_first = context.uow_ctx.checkpoints context.checkpoints_second = context.uow_ctx.checkpoints @then("both accesses should return the same CheckpointRepository instance") def step_checkpoints_cached(context: Context) -> None: """Verify both accesses return the same object (lazy caching).""" assert context.checkpoints_first is context.checkpoints_second, ( "Expected same CheckpointRepository instance on repeated access" ) # ========================================================================= # refresh() method (line 312) # ========================================================================= @when("I call refresh with a dummy entity") def step_call_refresh(context: Context) -> None: """Call refresh() with a dummy entity object.""" context.dummy_entity = MagicMock(name="DummyEntity") context.uow_ctx.refresh(context.dummy_entity) @then("session.refresh should have been called with that entity") def step_verify_session_refresh_called(context: Context) -> None: """Verify the underlying session.refresh was called with the entity.""" context.mock_session.refresh.assert_called_once_with(context.dummy_entity) # ========================================================================= # Non-SQLite engine options (lines 90-93 deeper verification) # ========================================================================= @given("a UnitOfWork configured with a postgresql-style database URL") def step_uow_postgresql_url(context: Context) -> None: """Create a UnitOfWork with a PostgreSQL URL.""" context.uow = UnitOfWork( database_url="postgresql+psycopg2://user:pass@db.example.com:5432/myapp" ) @when("I trigger engine creation with mocked dependencies") def step_trigger_engine_creation_mocked(context: Context) -> None: """Access engine property with mocked create_engine and migration runner.""" import sys as _sys from sqlalchemy.engine import Engine mock_engine = MagicMock(spec=Engine) # Ensure psycopg2 can be imported (it may not be installed in test env) _needs_psycopg2_stub = "psycopg2" not in _sys.modules if _needs_psycopg2_stub: _sys.modules["psycopg2"] = MagicMock() try: with ( patch.object(context.uow, "_ensure_database_initialized"), patch( "cleveragents.infrastructure.database.unit_of_work.create_engine", return_value=mock_engine, ) as mock_create, ): engine = context.uow.engine context.result_engine = engine context.mock_create_engine = mock_create finally: if _needs_psycopg2_stub: _sys.modules.pop("psycopg2", None) @then("create_engine should have been called without isolation_level") def step_create_engine_no_isolation_level(context: Context) -> None: """Verify create_engine was called without isolation_level kwarg.""" call_kwargs = context.mock_create_engine.call_args[1] assert "isolation_level" not in call_kwargs, ( f"isolation_level should not be set for non-SQLite, got {call_kwargs}" ) @then("create_engine should have been called with future=True") def step_create_engine_future_true(context: Context) -> None: """Verify create_engine was called with future=True.""" call_kwargs = context.mock_create_engine.call_args[1] assert call_kwargs.get("future") is True, f"Expected future=True, got {call_kwargs}" # ========================================================================= # In-memory engine cache hit (lines 70-78) # ========================================================================= @given("a first UnitOfWork has already created an in-memory engine") def step_first_uow_creates_engine(context: Context) -> None: """Create a UnitOfWork and access its engine, populating MEMORY_ENGINES.""" uow1 = UnitOfWork(database_url="sqlite:///:memory:") with patch.object(uow1, "_ensure_database_initialized"): engine1 = uow1.engine context.first_engine = engine1 @when("a second UnitOfWork accesses its engine property") def step_second_uow_accesses_engine(context: Context) -> None: """Create a second UnitOfWork and access its engine property.""" uow2 = UnitOfWork(database_url="sqlite:///:memory:") with patch.object(uow2, "_ensure_database_initialized"): engine2 = uow2.engine context.second_engine = engine2 @then("the second engine should be the same object as the first") def step_second_engine_same_as_first(context: Context) -> None: """Both UnitOfWork instances should share the same cached engine.""" from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES assert context.second_engine is context.first_engine, ( "Expected same engine object from MEMORY_ENGINES cache" ) # Restore cache if hasattr(context, "saved_memory_engines"): MEMORY_ENGINES.clear() MEMORY_ENGINES.update(context.saved_memory_engines) # ========================================================================= # Argument validation in __init__ (lines 68-75) # ========================================================================= @when("I try to create a UnitOfWork with an empty database_url") def step_uow_empty_database_url(context: Context) -> None: """Try to create UnitOfWork with empty database_url; capture ValueError.""" try: UnitOfWork(database_url="") context.last_error = None except ValueError as exc: context.last_error = exc @when("I try to create a UnitOfWork with pool_size set to {value:d}") def step_uow_invalid_pool_size(context: Context, value: int) -> None: """Try to create UnitOfWork with invalid pool_size; capture ValueError.""" try: UnitOfWork(database_url="sqlite:///test.db", pool_size=value) context.last_error = None except ValueError as exc: context.last_error = exc @when("I try to create a UnitOfWork with max_overflow set to {value:d}") def step_uow_invalid_max_overflow(context: Context, value: int) -> None: """Try to create UnitOfWork with invalid max_overflow; capture ValueError.""" try: UnitOfWork(database_url="sqlite:///test.db", max_overflow=value) context.last_error = None except ValueError as exc: context.last_error = exc @when("I try to create a UnitOfWork with pool_recycle set to {value:d}") def step_uow_invalid_pool_recycle(context: Context, value: int) -> None: """Try to create UnitOfWork with invalid pool_recycle; capture ValueError.""" try: UnitOfWork(database_url="sqlite:///test.db", pool_recycle=value) context.last_error = None except ValueError as exc: context.last_error = exc # ========================================================================= # dispose() method (lines 196-205) # ========================================================================= @given("a UnitOfWork with a mocked non-memory engine") def step_uow_mocked_non_memory_engine(context: Context) -> None: """Create a UnitOfWork with a pre-set mocked engine to test dispose().""" from sqlalchemy.engine import Engine context.uow = UnitOfWork( database_url="postgresql://user:pass@localhost:5432/testdb" ) context.mock_engine = MagicMock(spec=Engine) context.uow._engine = context.mock_engine @when("I call dispose on the UnitOfWork") def step_call_dispose_on_uow(context: Context) -> None: """Call dispose() on the UnitOfWork stored in context.""" context.uow.dispose() @then("the internal engine should have been disposed and cleared") def step_engine_disposed_and_cleared(context: Context) -> None: """Verify dispose() called engine.dispose() and cleared internal references.""" context.mock_engine.dispose.assert_called_once() assert context.uow._engine is None, ( f"Expected _engine to be None after dispose(), got {context.uow._engine}" ) assert context.uow._session_factory is None, ( "Expected _session_factory to be None after dispose()" ) @then("the dispose call should complete without error") def step_dispose_no_error(context: Context) -> None: """Verify dispose() completes without raising (False branch — engine is None).""" assert context.uow._engine is None, ( "Expected _engine to remain None when dispose() is a no-op" ) # ========================================================================= # Context manager protocol (__enter__ / __exit__) (lines 207-213) # ========================================================================= @when("I use the UnitOfWork as a context manager") def step_use_uow_as_context_manager(context: Context) -> None: """Use the UnitOfWork with a with-statement to exercise __enter__/__exit__.""" with context.uow as entered_uow: context.entered_uow = entered_uow @then("the engine should be disposed on context manager exit") def step_engine_disposed_on_context_exit(context: Context) -> None: """Verify __enter__ returned self and __exit__ triggered dispose().""" assert context.entered_uow is context.uow, "__enter__ should return self" context.mock_engine.dispose.assert_called_once() # ========================================================================= # psycopg2 ImportError hint (lines 142-148) # ========================================================================= @when("I access the engine with psycopg2 forcibly unavailable") def step_access_engine_no_psycopg2(context: Context) -> None: """Access engine property while psycopg2 is blocked via sys.modules.""" import sys as _sys try: with ( patch.dict(_sys.modules, {"psycopg2": None}), patch.object(context.uow, "_ensure_database_initialized"), ): _ = context.uow.engine context.raised_exception = None except ImportError as exc: context.raised_exception = exc @then("an ImportError should be raised mentioning the server extra") def step_import_error_server_extra(context: Context) -> None: """Verify ImportError was raised containing the install-hint message.""" assert context.raised_exception is not None, ( "Expected an ImportError to be raised but no exception was captured" ) assert isinstance(context.raised_exception, ImportError), ( f"Expected ImportError, got {type(context.raised_exception)}" ) assert "server" in str(context.raised_exception).lower(), ( f"Expected 'server' in exception message: '{context.raised_exception}'" )