mirror of
https://github.com/langgenius/dify.git
synced 2026-09-24 23:22:26 +08:00
test: use sqlite3 session in test_credit_pool_service (#38694)
This commit is contained in:
Asuka Minato
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user