routstr-core/tests/integration/test_child_keys.py

import asyncio
import secrets
from typing import Any

import pytest
from fastapi import HTTPException
from sqlmodel.ext.asyncio.session import AsyncSession

from routstr.auth import adjust_payment_for_tokens, pay_for_request
from routstr.balance import ChildKeyRequest, create_child_key
from routstr.core.db import ApiKey, create_session
from routstr.core.settings import settings


@pytest.mark.asyncio
async def test_child_key_flow(integration_session: AsyncSession) -> None:
    # 1. Create a parent key with balance
    parent_raw = "parent_test_key_" + secrets.token_hex(4)
    parent_key = ApiKey(
        hashed_key=parent_raw,
        balance=10000,  # 10 sats
    )
    integration_session.add(parent_key)
    await integration_session.commit()
    await integration_session.refresh(parent_key)

    # Mock settings
    settings.child_key_cost = 1000  # 1 sat

    # 2. Call create_child_key
    result = await create_child_key(
        ChildKeyRequest(count=1), parent_key, integration_session
    )

    assert "api_keys" in result
    assert result["cost_msats"] == 1000
    assert result["parent_balance"] == 9000

    child_key_raw = result["api_keys"][0][3:]  # remove sk-

    # 3. Verify child key exists in DB
    child_key_db = await integration_session.get(ApiKey, child_key_raw)
    assert child_key_db is not None
    assert child_key_db.parent_key_hash == parent_key.hashed_key
    assert child_key_db.balance == 0

    # 4. Test payment with child key
    cost = 500
    await pay_for_request(child_key_db, cost, integration_session)

    # Refresh keys
    await integration_session.refresh(parent_key)
    await integration_session.refresh(child_key_db)

    # Parent should be charged
    assert parent_key.reserved_balance == 500
    assert parent_key.total_requests == 1

    # Child should have total_requests incremented
    assert child_key_db.total_requests == 1

    # 5. Test adjustment
    response_data = {"model": "test-model", "usage": {"total_tokens": 10}}

    # Mock calculate_cost
    import routstr.auth
    from routstr.payment.cost_calculation import CostData

    async def mock_calculate_cost(*args: Any, **kwargs: Any) -> CostData:
        return CostData(
            base_msats=0, input_msats=200, output_msats=200, total_msats=400
        )

    # Patch calculate_cost
    original_calculate_cost = routstr.auth.calculate_cost
    routstr.auth.calculate_cost = mock_calculate_cost

    try:
        adjustment = await adjust_payment_for_tokens(
            child_key_db, response_data, integration_session, 500
        )
        assert adjustment["total_msats"] == 400

        # Refresh keys
        await integration_session.refresh(parent_key)
        await integration_session.refresh(child_key_db)

        # Parent should have updated balance and total_spent
        assert parent_key.reserved_balance == 0
        assert parent_key.balance == 9000 - 400
        assert (
            parent_key.total_spent == 1400
        )  # 1000 for child key creation + 400 for request

        # Child should also have total_spent updated
        assert child_key_db.total_spent == 400

    finally:
        routstr.auth.calculate_cost = original_calculate_cost


@pytest.mark.asyncio
async def test_child_key_insufficient_balance(
    integration_session: AsyncSession,
) -> None:
    parent_key = ApiKey(
        hashed_key="poor_parent_" + secrets.token_hex(4),
        balance=500,
    )
    integration_session.add(parent_key)
    await integration_session.commit()
    await integration_session.refresh(parent_key)

    settings.child_key_cost = 1000

    with pytest.raises(HTTPException) as exc:
        await create_child_key(
            ChildKeyRequest(count=1), parent_key, integration_session
        )
    assert exc.value.status_code == 402


@pytest.mark.asyncio
async def test_concurrent_child_key_creation_is_atomic(
    patched_db_engine: None,
) -> None:
    """Two concurrent create_child_key() calls with balance for exactly one must
    result in exactly one success and one 402, with the parent balance deducted
    only once."""
    child_key_cost = 1000
    settings.child_key_cost = child_key_cost

    parent_hash = f"parent_concurrent_{secrets.token_hex(8)}"
    async with create_session() as session:
        parent = ApiKey(hashed_key=parent_hash, balance=child_key_cost)
        session.add(parent)
        await session.commit()

    results: list[str] = []

    async def attempt() -> None:
        async with create_session() as session:
            fresh_parent = await session.get(ApiKey, parent_hash)
            assert fresh_parent is not None
            try:
                await create_child_key(ChildKeyRequest(count=1), fresh_parent, session)
                results.append("success")
            except HTTPException as exc:
                assert exc.status_code == 402
                results.append("blocked")

    await asyncio.gather(attempt(), attempt())

    assert sorted(results) == ["blocked", "success"], (
        f"Expected exactly one success and one 402, got: {results}"
    )

    async with create_session() as session:
        final = await session.get(ApiKey, parent_hash)
        assert final is not None

    assert final.balance == 0, (
        f"Balance should be fully deducted once: expected 0, got {final.balance}"
    )
    assert final.total_spent == child_key_cost, (
        f"total_spent should equal one deduction: expected {child_key_cost}, "
        f"got {final.total_spent}"
    )


@pytest.mark.asyncio
async def test_child_key_cannot_create_child(integration_session: AsyncSession) -> None:
    parent_key = ApiKey(
        hashed_key="parent_" + secrets.token_hex(4),
        balance=10000,
    )
    child_key = ApiKey(
        hashed_key="child_" + secrets.token_hex(4),
        balance=0,
        parent_key_hash=parent_key.hashed_key,
    )
    integration_session.add(parent_key)
    integration_session.add(child_key)
    await integration_session.commit()
    await integration_session.refresh(child_key)

    with pytest.raises(HTTPException) as exc:
        await create_child_key(ChildKeyRequest(count=1), child_key, integration_session)
    assert exc.value.status_code == 400
    assert "Cannot create a child key for another child key" in str(exc.value.detail)