diff --git a/routstr/auth.py b/routstr/auth.py index b1e16987..deb7b1d4 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -4,6 +4,7 @@ from typing import Optional from fastapi import HTTPException from sqlalchemy.exc import IntegrityError +from sqlalchemy import func from sqlmodel import col, update from .core import get_logger @@ -390,9 +391,11 @@ async def revert_pay_for_request( stmt = ( update(ApiKey) .where(col(ApiKey.hashed_key) == key.hashed_key) + .where(col(ApiKey.reserved_balance) >= cost_per_request) + .where(col(ApiKey.total_requests) > 0) .values( reserved_balance=col(ApiKey.reserved_balance) - cost_per_request, - total_requests=col(ApiKey.total_requests) - 1, + total_requests=func.greatest(col(ApiKey.total_requests) - 1, 0), ) ) diff --git a/tests/integration/test_reserved_balance_negative.py b/tests/integration/test_reserved_balance_negative.py index 21f9b6d0..457ca122 100644 --- a/tests/integration/test_reserved_balance_negative.py +++ b/tests/integration/test_reserved_balance_negative.py @@ -4,6 +4,7 @@ import asyncio import uuid import pytest +from fastapi import HTTPException from httpx import AsyncClient from sqlmodel.ext.asyncio.session import AsyncSession @@ -150,16 +151,12 @@ async def test_insufficient_reserved_balance_for_revert( await integration_session.commit() # Try to revert more than available - # Note: Current implementation allows reserved_balance to go negative - await revert_pay_for_request(test_key, integration_session, 100) + with pytest.raises(HTTPException) as excinfo: + await revert_pay_for_request(test_key, integration_session, 100) + assert excinfo.value.status_code == 402 + assert excinfo.value.detail["error"]["code"] == "payment_error" # Refresh to get updated values await integration_session.refresh(test_key) - - # Current implementation allows negative reserved balance - assert test_key.reserved_balance == -100, ( - f"Expected reserved_balance to be -100, got: {test_key.reserved_balance}" - ) - assert test_key.total_requests == -1, ( - f"Expected total_requests to be -1, got: {test_key.total_requests}" - ) + assert test_key.reserved_balance == 0 + assert test_key.total_requests == 0