"""Comprehensive database consistency tests""" import asyncio import time from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch import pytest from httpx import AsyncClient, Response from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlmodel import select from routstr.core.db import ApiKey class TestTransactionAtomicity: """Test transaction atomicity across all database operations""" @pytest.mark.asyncio async def test_balance_update_atomicity( self, authenticated_client: AsyncClient, integration_session: AsyncSession, db_snapshot: Any, ) -> None: """Test that balance updates are atomic and rolled back on failure""" # Get initial balance api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() initial_balance = api_key.balance # Test database atomicity by simulating a failed transaction # Create a new session for isolated transaction from sqlalchemy.ext.asyncio import AsyncSession async with AsyncSession(integration_session.bind) as test_session: try: # Get api key in new session result = await test_session.execute( select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] ) test_api_key = result.scalar_one() # Update balance test_api_key.balance -= 1000 await test_session.flush() # Apply changes but don't commit # Simulate an error that would cause rollback raise Exception("Simulated error after balance update") except Exception: await test_session.rollback() # Verify balance wasn't changed in main session await integration_session.refresh(api_key) assert api_key.balance == initial_balance # Test with concurrent modifications await db_snapshot.capture() # Try to update in a transaction that will fail from sqlalchemy import update try: await integration_session.execute( update(ApiKey) .where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] .values(balance=ApiKey.balance - 1000) ) # Force a constraint violation or error await integration_session.execute( update(ApiKey) .where(ApiKey.hashed_key == "non_existent_key") # type: ignore[arg-type] .values(balance=-1) # This should fail ) await integration_session.commit() except Exception: await integration_session.rollback() # Verify no changes were persisted diff = await db_snapshot.diff() assert len(diff["api_keys"]["added"]) == 0 assert len(diff["api_keys"]["modified"]) == 0 @pytest.mark.asyncio async def test_topup_rollback_on_failure( self, authenticated_client: AsyncClient, integration_session: AsyncSession, db_snapshot: Any, ) -> None: """Test that failed top-ups don't leave partial database state""" # Get initial state api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() initial_balance = api_key.balance # Mock wallet to fail after token validation with patch("routstr.wallet.send_token") as mock_wallet_func: mock_proof = MagicMock() mock_proof.amount = 1000 mock_wallet = AsyncMock() mock_wallet.deserialize_token = AsyncMock(return_value=[mock_proof]) mock_wallet.redeem = AsyncMock( side_effect=Exception("Network error during redemption") ) mock_wallet_func.return_value = mock_wallet # Attempt top-up response = await authenticated_client.post( "/v1/wallet/topup", params={"cashu_token": "cashuAey..."} ) # The mock returns 400 for invalid tokens assert response.status_code in [400, 500] # Verify no balance change await integration_session.refresh(api_key) assert api_key.balance == initial_balance # Verify clean database state diff = await db_snapshot.diff() assert len(diff["api_keys"]["added"]) == 0 assert len(diff["api_keys"]["modified"]) == 0 @pytest.mark.asyncio async def test_concurrent_balance_updates( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test atomic balance updates under concurrent operations""" # Get API key info api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) # Set a known balance stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() api_key.balance = 10000 await integration_session.commit() # Simulate concurrent balance updates through direct database operations async def update_balance(session: AsyncSession, amount: int) -> bool: stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await session.execute(stmt) key = result.scalar_one() key.balance -= amount key.total_spent += amount key.total_requests += 1 try: await session.commit() return True except Exception: await session.rollback() return False # Run concurrent balance updates tasks = [] deduction_amounts = [100, 200, 300, 400, 500] for amount in deduction_amounts: # Create a new session for each concurrent operation async with AsyncSession(integration_session.bind) as session: task = update_balance(session, amount) tasks.append(task) await asyncio.gather(*tasks, return_exceptions=True) # Verify final balance is consistent await integration_session.refresh(api_key) # Balance should have some deduction but exact amount depends on implementation assert api_key.balance < 10000 assert api_key.balance >= 0 # Should never go negative class TestConcurrentOperations: """Test database consistency under concurrent operations""" @pytest.mark.asyncio async def test_multiple_requests_same_api_key( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test multiple concurrent requests with the same API key""" # Mock the wallet info endpoint to track concurrent calls call_count = 0 call_times = [] async def track_concurrent_calls() -> Dict[str, int]: nonlocal call_count call_count += 1 call_times.append(time.time()) await asyncio.sleep(0.1) # Simulate processing time return {"balance": 1000} # Make 10 concurrent requests tasks = [] for _ in range(10): task = authenticated_client.get("/v1/wallet/info") tasks.append(task) responses = await asyncio.gather(*tasks) # All requests should succeed for response in responses: assert response.status_code == 200 @pytest.mark.asyncio async def test_simultaneous_topup_and_usage( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test simultaneous top-up and balance usage operations""" # Get API key info api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) # Set initial balance stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() initial_balance = 5000 api_key.balance = initial_balance await integration_session.commit() # Mock wallet for topup with patch("routstr.wallet.send_token") as mock_wallet_func: mock_proof = MagicMock() mock_proof.amount = 2000 mock_wallet = AsyncMock() mock_wallet.deserialize_token = AsyncMock(return_value=[mock_proof]) mock_wallet.redeem = AsyncMock(return_value=[mock_proof]) mock_wallet_func.return_value = mock_wallet # Mock proxy endpoint to simulate usage with patch("httpx.AsyncClient.request") as mock_request: # Mock successful proxy response mock_response = AsyncMock() mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} mock_response.aiter_bytes = AsyncMock( return_value=iter([b'{"result": "ok"}']) ) mock_response.is_stream_consumed = False mock_request.return_value = mock_response # Run topup and usage concurrently async def topup() -> Any: return await authenticated_client.post( "/v1/wallet/topup", params={"cashu_token": "cashuAey..."} ) async def use_balance() -> Any: # This would normally deduct balance return await authenticated_client.post( "/v1/chat/completions", json={"model": "test", "messages": []} ) # Execute concurrently results = await asyncio.gather( topup(), use_balance(), return_exceptions=True ) topup_result = results[0] usage_result = results[1] # At least one should succeed assert not isinstance(topup_result, Exception) or not isinstance( usage_result, Exception ) # Verify final balance is consistent await integration_session.refresh(api_key) # Balance should be between initial and initial + topup amount assert api_key.balance >= initial_balance assert api_key.balance <= initial_balance + 2000 @pytest.mark.asyncio async def test_race_condition_prevention( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test that race conditions are prevented in balance updates""" # Get API key info api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) # Set a specific balance stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() api_key.balance = 1000 api_key.total_spent = 0 api_key.total_requests = 0 await integration_session.commit() # Create a controlled race condition scenario balance_checks: List[int] = [] async def check_and_update_balance() -> bool: # Read current balance stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) current_api_key = result.scalar_one() current_balance = current_api_key.balance balance_checks.append(current_balance) # Simulate processing delay await asyncio.sleep(0.01) # Try to update based on read value current_api_key.balance = current_balance - 100 current_api_key.total_spent += 100 current_api_key.total_requests += 1 try: await integration_session.commit() return True except Exception: await integration_session.rollback() return False # Run multiple concurrent updates tasks = [check_and_update_balance() for _ in range(5)] results = await asyncio.gather(*tasks, return_exceptions=True) # Refresh and check final state await integration_session.refresh(api_key) # At least some updates should succeed successful_updates = sum(1 for r in results if r is True) assert successful_updates > 0 # Final balance should reflect successful updates expected_balance = 1000 - (successful_updates * 100) assert api_key.balance == expected_balance assert api_key.total_spent == successful_updates * 100 assert api_key.total_requests == successful_updates class TestDataIntegrity: """Test data integrity constraints and validations""" @pytest.mark.asyncio @pytest.mark.skip(reason="Balance never negative is not implemented") async def test_balance_never_negative( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test that balance can never go negative""" # Get API key info api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) # Set low balance stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() api_key.balance = 0 await integration_session.commit() # Try to refund more than balance response = await authenticated_client.post( "/v1/wallet/refund", json={"amount": 1000} ) # Should fail assert response.status_code == 400 assert "Balance too small to refund" in response.json()["detail"] # Verify balance unchanged await integration_session.refresh(api_key) assert api_key.balance == 100 @pytest.mark.asyncio async def test_primary_key_uniqueness( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test that primary key constraints are enforced""" # Get existing API key hash from authenticated client api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) # Try to manually insert duplicate key with same hash duplicate_key = ApiKey( hashed_key=api_key_hash, balance=5000, total_spent=0, total_requests=0 ) integration_session.add(duplicate_key) # Should raise integrity error with pytest.raises(IntegrityError): await integration_session.commit() await integration_session.rollback() @pytest.mark.asyncio async def test_timestamp_consistency( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test that timestamps are consistent and properly ordered""" # Track request times request_times: List[float] = [] # Make several requests with delays for i in range(3): start_time = time.time() response = await authenticated_client.get("/v1/wallet/info") assert response.status_code == 200 request_times.append(start_time) await asyncio.sleep(0.1) # Verify timestamps are monotonically increasing for i in range(1, len(request_times)): assert request_times[i] > request_times[i - 1] @pytest.mark.asyncio async def test_numeric_field_constraints( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test constraints on numeric fields""" # Get API key api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() # Test setting invalid values directly # These should maintain integrity assert api_key.balance >= 0 assert api_key.total_spent >= 0 assert api_key.total_requests >= 0 # Verify calculations are consistent if api_key.total_requests > 0: average_cost = api_key.total_spent / api_key.total_requests assert average_cost >= 0 class TestPerformance: """Test database performance characteristics""" @pytest.mark.asyncio async def test_operation_latency( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test that database operations complete within acceptable time""" operation_times: Dict[str, List[float]] = { "select": [], "update": [], "insert": [], } # Test SELECT performance for _ in range(10): start = time.time() response = await authenticated_client.get("/v1/wallet/info") end = time.time() assert response.status_code == 200 operation_times["select"].append((end - start) * 1000) # Convert to ms # Test UPDATE performance (via topup) with patch("routstr.wallet.send_token") as mock_wallet_func: mock_proof = MagicMock() mock_proof.amount = 100 mock_wallet = AsyncMock() mock_wallet.deserialize_token = AsyncMock(return_value=[mock_proof]) mock_wallet.redeem = AsyncMock(return_value=[mock_proof]) mock_wallet_func.return_value = mock_wallet for _ in range(5): start = time.time() response = await authenticated_client.post( "/v1/wallet/topup", params={"cashu_token": "cashuAey..."} ) end = time.time() # Skip if token is invalid (400) if response.status_code == 400: continue assert response.status_code == 200 operation_times["update"].append((end - start) * 1000) # Verify all operations < 100ms for op_type, times in operation_times.items(): if times: # Only check if we have measurements avg_time = sum(times) / len(times) max_time = max(times) # Average should be well under 100ms assert avg_time < 100, ( f"{op_type} average time {avg_time}ms exceeds 100ms" ) # No single operation should exceed 200ms assert max_time < 200, f"{op_type} max time {max_time}ms exceeds 200ms" @pytest.mark.asyncio async def test_connection_pool_behavior( self, authenticated_client: AsyncClient, integration_app: Any, ) -> None: """Test database connection pool behavior under load""" # Make many concurrent requests to test connection pooling async def make_request() -> Response: return await authenticated_client.get("/v1/wallet/info") # Create 50 concurrent requests tasks = [make_request() for _ in range(50)] start = time.time() responses = await asyncio.gather(*tasks, return_exceptions=True) end = time.time() # All should succeed success_count = sum( 1 for r in responses if not isinstance(r, Exception) and hasattr(r, "status_code") and r.status_code == 200 ) assert success_count == 50, f"Only {success_count}/50 requests succeeded" # Should complete reasonably quickly (< 5 seconds for 50 requests) total_time = end - start assert total_time < 5.0, f"50 concurrent requests took {total_time}s" @pytest.mark.asyncio async def test_index_usage( self, authenticated_client: AsyncClient, integration_session: AsyncSession, ) -> None: """Test that database indexes are used efficiently""" # Get API key for testing api_key_header = authenticated_client.headers["Authorization"].replace( "Bearer ", "" ) # API key format is "sk-{hashed_key}" api_key_hash = ( api_key_header[3:] if api_key_header.startswith("sk-") else api_key_header ) # Primary key lookup should be fast start = time.time() stmt = select(ApiKey).where(ApiKey.hashed_key == api_key_hash) # type: ignore[arg-type] result = await integration_session.execute(stmt) api_key = result.scalar_one() end = time.time() lookup_time = (end - start) * 1000 assert lookup_time < 10, f"Primary key lookup took {lookup_time}ms" # Verify we got the right record assert api_key.hashed_key == api_key_hash