diff --git a/api/tests/unit_tests/services/test_credit_pool_service.py b/api/tests/unit_tests/services/test_credit_pool_service.py index caf5d6ade02..f3b97e8bbd6 100644 --- a/api/tests/unit_tests/services/test_credit_pool_service.py +++ b/api/tests/unit_tests/services/test_credit_pool_service.py @@ -1,12 +1,13 @@ +"""Credit-pool accounting tests backed by real SQLite sessions.""" + from collections.abc import Generator from types import SimpleNamespace from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 import pytest -from sqlalchemy import create_engine, select -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy import select +from sqlalchemy.orm import Session from core.errors.error import QuotaExceededError from models import TenantCreditPool @@ -20,32 +21,20 @@ from services.credit_pool_service import ( ) -def _create_engine_with_pool(*, quota_limit: int, quota_used: int) -> tuple[Engine, str, str]: - engine = create_engine("sqlite:///:memory:") - TenantCreditPool.__table__.create(engine) - tenant_id = str(uuid4()) - pool_id = str(uuid4()) - with engine.begin() as connection: - connection.execute( - TenantCreditPool.__table__.insert(), - { - "id": pool_id, - "tenant_id": tenant_id, - "pool_type": ProviderQuotaType.TRIAL, - "quota_limit": quota_limit, - "quota_used": quota_used, - }, - ) - return engine, tenant_id, pool_id +def _create_pool(session: Session, *, quota_limit: int, quota_used: int) -> TenantCreditPool: + pool = TenantCreditPool( + tenant_id=str(uuid4()), + pool_type=ProviderQuotaType.TRIAL, + quota_limit=quota_limit, + quota_used=quota_used, + ) + session.add(pool) + session.commit() + return pool -def _make_session(engine: Engine) -> Session: - return sessionmaker(bind=engine, expire_on_commit=False)() - - -def _get_quota_used(*, engine: Engine, pool_id: str) -> int | None: - with engine.connect() as connection: - return connection.scalar(select(TenantCreditPool.quota_used).where(TenantCreditPool.id == pool_id)) +def _get_quota_used(*, session: Session, pool_id: str) -> int | None: + return session.scalar(select(TenantCreditPool.quota_used).where(TenantCreditPool.id == pool_id)) def _make_redis_lock() -> MagicMock: @@ -60,14 +49,16 @@ def _disable_billing_quota_by_default() -> Generator[None, None, None]: yield -def test_get_pool_uses_provided_session() -> None: - engine, tenant_id, _ = _create_engine_with_pool(quota_limit=10, quota_used=2) - - with _make_session(engine) as session: - pool = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL, session=session) +def test_get_pool_uses_provided_session(sqlite_session: Session) -> None: + persisted_pool = _create_pool(sqlite_session, quota_limit=10, quota_used=2) + pool = CreditPoolService.get_pool( + tenant_id=persisted_pool.tenant_id, + pool_type=ProviderQuotaType.TRIAL, + session=sqlite_session, + ) assert pool is not None - assert pool.tenant_id == tenant_id + assert pool.tenant_id == persisted_pool.tenant_id assert pool.quota_used == 2 @@ -78,132 +69,127 @@ def test_credit_pool_balance_unlimited_remaining_and_sufficiency() -> None: assert pool.has_sufficient_credits(10_000) -def test_check_and_deduct_credits_deducts_exact_amount_when_sufficient() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) +def test_check_and_deduct_credits_deducts_exact_amount_when_sufficient(sqlite_session: Session) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=2) - with _make_session(engine) as session: - deducted_credits = CreditPoolService.check_and_deduct_credits( - tenant_id=tenant_id, credits_required=3, session=session - ) + deducted_credits = CreditPoolService.check_and_deduct_credits( + tenant_id=pool.tenant_id, credits_required=3, session=sqlite_session + ) assert deducted_credits == 3 - assert _get_quota_used(engine=engine, pool_id=pool_id) == 5 + assert sqlite_session.in_transaction() is False + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 5 -def test_check_and_deduct_credits_returns_zero_for_non_positive_request() -> None: +def test_check_and_deduct_credits_returns_zero_for_non_positive_request(sqlite_session: Session) -> None: assert ( - CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=0, session=MagicMock()) == 0 + CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=0, session=sqlite_session) + == 0 ) -def test_check_and_deduct_credits_raises_when_pool_is_missing() -> None: - engine = create_engine("sqlite:///:memory:") - TenantCreditPool.__table__.create(engine) - - with _make_session(engine) as session, pytest.raises(QuotaExceededError, match="Credit pool not found"): - CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=1, session=session) +def test_check_and_deduct_credits_raises_when_pool_is_missing(sqlite_session: Session) -> None: + with pytest.raises(QuotaExceededError, match="Credit pool not found"): + CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=1, session=sqlite_session) -def test_check_and_deduct_credits_raises_when_pool_is_empty() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=10) +def test_check_and_deduct_credits_raises_when_pool_is_empty(sqlite_session: Session) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=10) - with _make_session(engine) as session, pytest.raises(QuotaExceededError, match="No credits remaining"): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1, session=session) + with pytest.raises(QuotaExceededError, match="No credits remaining"): + CreditPoolService.check_and_deduct_credits(tenant_id=pool.tenant_id, credits_required=1, session=sqlite_session) - assert _get_quota_used(engine=engine, pool_id=pool_id) == 10 + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 10 -def test_check_and_deduct_credits_raises_without_partial_deduction_when_insufficient() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=9) +def test_check_and_deduct_credits_raises_without_partial_deduction_when_insufficient( + sqlite_session: Session, +) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=9) - with _make_session(engine) as session, pytest.raises(QuotaExceededError, match="Insufficient credits remaining"): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=3, session=session) + with pytest.raises(QuotaExceededError, match="Insufficient credits remaining"): + CreditPoolService.check_and_deduct_credits(tenant_id=pool.tenant_id, credits_required=3, session=sqlite_session) - assert _get_quota_used(engine=engine, pool_id=pool_id) == 9 + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 9 -def test_check_and_deduct_credits_wraps_unexpected_deduction_errors() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) +def test_check_and_deduct_credits_wraps_unexpected_deduction_errors(sqlite_session: Session) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=2) with ( - _make_session(engine) as session, patch.object(CreditPoolService, "_get_locked_pool", side_effect=RuntimeError("database unavailable")), pytest.raises(QuotaExceededError, match="Failed to deduct credits"), ): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1, session=session) + CreditPoolService.check_and_deduct_credits(tenant_id=pool.tenant_id, credits_required=1, session=sqlite_session) - assert _get_quota_used(engine=engine, pool_id=pool_id) == 2 + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 2 -def test_deduct_credits_capped_returns_zero_for_non_positive_request() -> None: - assert CreditPoolService.deduct_credits_capped(tenant_id=str(uuid4()), credits_required=0, session=MagicMock()) == 0 +def test_deduct_credits_capped_returns_zero_for_non_positive_request(sqlite_session: Session) -> None: + assert ( + CreditPoolService.deduct_credits_capped(tenant_id=str(uuid4()), credits_required=0, session=sqlite_session) == 0 + ) -def test_deduct_credits_capped_returns_zero_when_pool_is_missing() -> None: - engine = create_engine("sqlite:///:memory:") - TenantCreditPool.__table__.create(engine) - - with _make_session(engine) as session: - deducted_credits = CreditPoolService.deduct_credits_capped( - tenant_id=str(uuid4()), credits_required=1, session=session - ) +def test_deduct_credits_capped_returns_zero_when_pool_is_missing(sqlite_session: Session) -> None: + deducted_credits = CreditPoolService.deduct_credits_capped( + tenant_id=str(uuid4()), credits_required=1, session=sqlite_session + ) assert deducted_credits == 0 -def test_deduct_credits_capped_returns_zero_when_pool_is_empty() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=10) +def test_deduct_credits_capped_returns_zero_when_pool_is_empty(sqlite_session: Session) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=10) - with _make_session(engine) as session: - deducted_credits = CreditPoolService.deduct_credits_capped( - tenant_id=tenant_id, credits_required=1, session=session - ) + deducted_credits = CreditPoolService.deduct_credits_capped( + tenant_id=pool.tenant_id, credits_required=1, session=sqlite_session + ) assert deducted_credits == 0 - assert _get_quota_used(engine=engine, pool_id=pool_id) == 10 + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 10 -def test_deduct_credits_capped_deducts_only_remaining_balance_when_insufficient() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=9) +def test_deduct_credits_capped_deducts_only_remaining_balance_when_insufficient( + sqlite_session: Session, +) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=9) - with _make_session(engine) as session: - deducted_credits = CreditPoolService.deduct_credits_capped( - tenant_id=tenant_id, credits_required=3, session=session - ) + deducted_credits = CreditPoolService.deduct_credits_capped( + tenant_id=pool.tenant_id, credits_required=3, session=sqlite_session + ) assert deducted_credits == 1 - assert _get_quota_used(engine=engine, pool_id=pool_id) == 10 + assert sqlite_session.in_transaction() is False + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 10 -def test_deduct_credits_capped_wraps_unexpected_deduction_errors() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) +def test_deduct_credits_capped_wraps_unexpected_deduction_errors(sqlite_session: Session) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=2) with ( - _make_session(engine) as session, patch.object(CreditPoolService, "_get_locked_pool", side_effect=RuntimeError("database unavailable")), pytest.raises(QuotaExceededError, match="Failed to deduct credits"), ): - CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1, session=session) + CreditPoolService.deduct_credits_capped(tenant_id=pool.tenant_id, credits_required=1, session=sqlite_session) - assert _get_quota_used(engine=engine, pool_id=pool_id) == 2 + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 2 -def test_deduct_credits_capped_reraises_quota_exceeded_errors() -> None: - engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) +def test_deduct_credits_capped_reraises_quota_exceeded_errors(sqlite_session: Session) -> None: + pool = _create_pool(sqlite_session, quota_limit=10, quota_used=2) with ( - _make_session(engine) as session, patch.object(CreditPoolService, "_get_locked_pool", side_effect=QuotaExceededError("quota unavailable")), pytest.raises(QuotaExceededError, match="quota unavailable"), ): - CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1, session=session) + CreditPoolService.deduct_credits_capped(tenant_id=pool.tenant_id, credits_required=1, session=sqlite_session) - assert _get_quota_used(engine=engine, pool_id=pool_id) == 2 + assert _get_quota_used(session=sqlite_session, pool_id=pool.id) == 2 -def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() -> None: +def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction(sqlite_session: Session) -> None: tenant_id = "tenant-1" - session = MagicMock() pool = SimpleNamespace(remaining_credits=10, quota_used=2) redis_lock = _make_redis_lock() @@ -215,7 +201,7 @@ def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() - tenant_id=tenant_id, credits_required=3, pool_type=ProviderQuotaType.TRIAL, - session=session, + session=sqlite_session, ) assert result == 3 @@ -227,12 +213,11 @@ def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() - ) redis_lock.acquire.assert_called_once_with(blocking=True) redis_lock.release.assert_called_once_with() - get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="trial") + get_locked_pool.assert_called_once_with(session=sqlite_session, tenant_id=tenant_id, pool_type="trial") -def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> None: +def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction(sqlite_session: Session) -> None: tenant_id = "tenant-1" - session = MagicMock() pool = SimpleNamespace(remaining_credits=2, quota_used=8) redis_lock = _make_redis_lock() @@ -244,7 +229,7 @@ def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> N tenant_id=tenant_id, credits_required=5, pool_type=ProviderQuotaType.PAID, - session=session, + session=sqlite_session, ) assert result == 2 @@ -256,7 +241,7 @@ def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> N ) redis_lock.acquire.assert_called_once_with(blocking=True) redis_lock.release.assert_called_once_with() - get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="paid") + get_locked_pool.assert_called_once_with(session=sqlite_session, tenant_id=tenant_id, pool_type="paid") def test_get_pool_uses_billing_quota_balance_when_enabled() -> None: @@ -419,28 +404,28 @@ def test_deduct_credits_capped_uses_billing_consume_capped_when_enabled() -> Non CreditPoolService.deduct_credits_capped, ], ) -def test_non_positive_credit_request_skips_tenant_redis_lock(deduct_method) -> None: +def test_non_positive_credit_request_skips_tenant_redis_lock( + deduct_method, + sqlite_session: Session, +) -> None: with patch("services.credit_pool_service.redis_client.lock") as lock: - result = deduct_method(tenant_id="tenant-1", credits_required=0, session=MagicMock()) + result = deduct_method(tenant_id="tenant-1", credits_required=0, session=sqlite_session) assert result == 0 lock.assert_not_called() -def test_check_and_deduct_credits_wraps_redis_lock_errors_without_querying_db() -> None: - session = MagicMock() +def test_check_and_deduct_credits_wraps_redis_lock_errors_without_querying_db(sqlite_session: Session) -> None: + with patch("services.credit_pool_service.redis_client.lock", side_effect=RuntimeError("redis unavailable")): + with pytest.raises(QuotaExceededError, match="Failed to deduct credits"): + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=1, session=sqlite_session) - with ( - patch("services.credit_pool_service.redis_client.lock", side_effect=RuntimeError("redis unavailable")), - pytest.raises(QuotaExceededError, match="Failed to deduct credits"), - ): - CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=1, session=session) - - session.scalar.assert_not_called() + assert sqlite_session.in_transaction() is False -def test_deduct_credits_capped_ignores_release_errors_after_successful_deduction() -> None: - session = MagicMock() +def test_deduct_credits_capped_ignores_release_errors_after_successful_deduction( + sqlite_session: Session, +) -> None: pool = SimpleNamespace(remaining_credits=3, quota_used=7) redis_lock = _make_redis_lock() redis_lock.release.side_effect = RuntimeError("release failed") @@ -449,7 +434,9 @@ def test_deduct_credits_capped_ignores_release_errors_after_successful_deduction patch("services.credit_pool_service.redis_client.lock", return_value=redis_lock), patch.object(CreditPoolService, "_get_locked_pool", return_value=pool), ): - result = CreditPoolService.deduct_credits_capped(tenant_id="tenant-1", credits_required=2, session=session) + result = CreditPoolService.deduct_credits_capped( + tenant_id="tenant-1", credits_required=2, session=sqlite_session + ) assert result == 2 assert pool.quota_used == 9