From daf17f51ab79d60fda29dfb33f18a67a8571e357 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sat, 10 Jan 2026 20:00:52 +0100 Subject: [PATCH] fix: move child-key route before catch-all and fix indentation --- examples/create_child_keys.py | 44 +++++++++ routstr/balance.py | 24 ++--- routstr/core/main.py | 8 +- tests/integration/test_child_keys.py | 131 +++++++++++++++++++++++++++ 4 files changed, 190 insertions(+), 17 deletions(-) create mode 100644 examples/create_child_keys.py create mode 100644 tests/integration/test_child_keys.py diff --git a/examples/create_child_keys.py b/examples/create_child_keys.py new file mode 100644 index 00000000..eefa3710 --- /dev/null +++ b/examples/create_child_keys.py @@ -0,0 +1,44 @@ +import httpx +import sys +import json + + +def create_child_keys(base_url, api_key, count=3): + headers = {"Authorization": f"Bearer {api_key}"} + + print(f"Requesting {count} child keys from {base_url}...") + + child_keys = [] + + for i in range(count): + try: + response = httpx.post(f"{base_url}/v1/balance/child-key", headers=headers) + if response.status_code == 200: + data = response.json() + child_keys.append(data["api_key"]) + print( + f" [{i + 1}] Created: {data['api_key']} (Cost: {data['cost_msats']} msats)" + ) + else: + print(f" [{i + 1}] Failed: {response.status_code} - {response.text}") + except Exception as e: + print(f" [{i + 1}] Error: {str(e)}") + + return child_keys + + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python create_child_keys.py [base_url]") + sys.exit(1) + + auth_key = sys.argv[1] + base_url = sys.argv[2] if len(sys.argv) > 2 else "http://localhost:8000" + + keys = create_child_keys(base_url, auth_key) + + if keys: + print("\nSuccessfully created child keys:") + print(json.dumps(keys, indent=2)) + else: + print("\nNo child keys were created.") diff --git a/routstr/balance.py b/routstr/balance.py index 472f8b05..3ee8ca11 100644 --- a/routstr/balance.py +++ b/routstr/balance.py @@ -251,18 +251,6 @@ async def donate(token: str, ref: str | None = None) -> str: return "Invalid token." -@router.api_route( - "/{path:path}", - methods=["GET", "POST", "PUT", "DELETE"], - include_in_schema=False, - response_model=None, -) -async def wallet_catch_all(path: str) -> NoReturn: - raise HTTPException( - status_code=404, detail="Not found check /docs for available endpoints" - ) - - @router.post("/child-key") async def create_child_key( key: ApiKey = Depends(get_key_from_header), @@ -310,6 +298,18 @@ async def create_child_key( } +@router.api_route( + "/{path:path}", + methods=["GET", "POST", "PUT", "DELETE"], + include_in_schema=False, + response_model=None, +) +async def wallet_catch_all(path: str) -> NoReturn: + raise HTTPException( + status_code=404, detail="Not found check /docs for available endpoints" + ) + + balance_router.include_router(lightning_router) balance_router.include_router(router) diff --git a/routstr/core/main.py b/routstr/core/main.py index 5bce5d32..3083097d 100644 --- a/routstr/core/main.py +++ b/routstr/core/main.py @@ -13,12 +13,10 @@ from starlette.exceptions import HTTPException from ..balance import balance_router, deprecated_wallet_router from ..discovery import providers_cache_refresher, providers_router from ..nip91 import announce_provider -from ..payment.models import ( - models_router, - update_sats_pricing, -) +from ..payment.models import models_router, update_sats_pricing from ..payment.price import update_prices_periodically -from ..proxy import initialize_upstreams, proxy_router, refresh_model_maps_periodically +from ..proxy import (initialize_upstreams, proxy_router, + refresh_model_maps_periodically) from ..wallet import periodic_payout from .admin import admin_router from .db import create_session, init_db, run_migrations diff --git a/tests/integration/test_child_keys.py b/tests/integration/test_child_keys.py new file mode 100644 index 00000000..77da7675 --- /dev/null +++ b/tests/integration/test_child_keys.py @@ -0,0 +1,131 @@ +import pytest +import secrets +from fastapi import HTTPException +from routstr.core.db import ApiKey +from routstr.balance import create_child_key +from routstr.auth import pay_for_request, adjust_payment_for_tokens +from routstr.core.settings import settings + + +@pytest.mark.asyncio +async def test_child_key_flow(integration_session): + # 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(parent_key, integration_session) + + assert "api_key" in result + assert result["cost_msats"] == 1000 + assert result["parent_balance"] == 9000 + + child_key_raw = result["api_key"][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, **kwargs): + 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): + 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(parent_key, integration_session) + assert exc.value.status_code == 402 + + +@pytest.mark.asyncio +async def test_child_key_cannot_create_child(integration_session): + 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(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)