import pytest from decimal import Decimal from datetime import date, datetime from httpx import AsyncClient from fastapi import status from tortoise.contrib.test import initializer, finalizer from rbox.main import app from rbox.models import User from rbox.billing.models import PricingConfig, Invoice, UsageAggregate, UserSubscription from rbox.auth import create_access_token @pytest.fixture(scope="module") def event_loop(): import asyncio loop = asyncio.get_event_loop_policy().new_event_loop() yield loop loop.close() @pytest.fixture(scope="module", autouse=True) async def initialize_tests(): initializer(["rbox.models", "rbox.billing.models"], db_url="sqlite://:memory:") yield await finalizer() @pytest.fixture async def test_user(): user = await User.create( username="testuser", email="test@example.com", hashed_password="hashed_password_here", is_active=True, is_superuser=False ) yield user await user.delete() @pytest.fixture async def admin_user(): user = await User.create( username="adminuser", email="admin@example.com", hashed_password="hashed_password_here", is_active=True, is_superuser=True ) yield user await user.delete() @pytest.fixture async def auth_token(test_user): token = create_access_token(data={"sub": test_user.username}) return token @pytest.fixture async def admin_token(admin_user): token = create_access_token(data={"sub": admin_user.username}) return token @pytest.fixture async def pricing_config(): configs = [] configs.append(await PricingConfig.create( config_key="storage_per_gb_month", config_value=Decimal("0.0045"), description="Storage cost per GB per month", unit="per_gb_month" )) configs.append(await PricingConfig.create( config_key="bandwidth_egress_per_gb", config_value=Decimal("0.009"), description="Bandwidth egress cost per GB", unit="per_gb" )) configs.append(await PricingConfig.create( config_key="free_tier_storage_gb", config_value=Decimal("15"), description="Free tier storage in GB", unit="gb" )) configs.append(await PricingConfig.create( config_key="free_tier_bandwidth_gb", config_value=Decimal("15"), description="Free tier bandwidth in GB per month", unit="gb" )) yield configs for config in configs: await config.delete() @pytest.mark.asyncio async def test_get_current_usage(test_user, auth_token): async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( "/api/billing/usage/current", headers={"Authorization": f"Bearer {auth_token}"} ) assert response.status_code == status.HTTP_200_OK data = response.json() assert "storage_gb" in data assert "bandwidth_down_gb_today" in data @pytest.mark.asyncio async def test_get_monthly_usage(test_user, auth_token): today = date.today() await UsageAggregate.create( user=test_user, date=today, storage_bytes_avg=1024 ** 3 * 10, storage_bytes_peak=1024 ** 3 * 12, bandwidth_up_bytes=1024 ** 3 * 2, bandwidth_down_bytes=1024 ** 3 * 5 ) async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( f"/api/billing/usage/monthly?year={today.year}&month={today.month}", headers={"Authorization": f"Bearer {auth_token}"} ) assert response.status_code == status.HTTP_200_OK data = response.json() assert data["storage_gb_avg"] == pytest.approx(10.0, rel=0.01) await UsageAggregate.filter(user=test_user).delete() @pytest.mark.asyncio async def test_get_subscription(test_user, auth_token): async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( "/api/billing/subscription", headers={"Authorization": f"Bearer {auth_token}"} ) assert response.status_code == status.HTTP_200_OK data = response.json() assert data["billing_type"] == "pay_as_you_go" assert data["status"] == "active" await UserSubscription.filter(user=test_user).delete() @pytest.mark.asyncio async def test_list_invoices(test_user, auth_token): invoice = await Invoice.create( user=test_user, invoice_number="INV-000001-202311", period_start=date(2023, 11, 1), period_end=date(2023, 11, 30), subtotal=Decimal("10.00"), tax=Decimal("0.00"), total=Decimal("10.00"), status="open" ) async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( "/api/billing/invoices", headers={"Authorization": f"Bearer {auth_token}"} ) assert response.status_code == status.HTTP_200_OK data = response.json() assert len(data) > 0 assert data[0]["invoice_number"] == "INV-000001-202311" await invoice.delete() @pytest.mark.asyncio async def test_get_invoice(test_user, auth_token): invoice = await Invoice.create( user=test_user, invoice_number="INV-000002-202311", period_start=date(2023, 11, 1), period_end=date(2023, 11, 30), subtotal=Decimal("10.00"), tax=Decimal("0.00"), total=Decimal("10.00"), status="open" ) async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( f"/api/billing/invoices/{invoice.id}", headers={"Authorization": f"Bearer {auth_token}"} ) assert response.status_code == status.HTTP_200_OK data = response.json() assert data["invoice_number"] == "INV-000002-202311" await invoice.delete() @pytest.mark.asyncio async def test_get_pricing(): async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get("/api/billing/pricing") assert response.status_code == status.HTTP_200_OK data = response.json() assert isinstance(data, dict) @pytest.mark.asyncio async def test_admin_get_pricing(admin_user, admin_token, pricing_config): async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( "/api/admin/billing/pricing", headers={"Authorization": f"Bearer {admin_token}"} ) assert response.status_code == status.HTTP_200_OK data = response.json() assert len(data) > 0 @pytest.mark.asyncio async def test_admin_update_pricing(admin_user, admin_token, pricing_config): config_id = pricing_config[0].id async with AsyncClient(app=app, base_url="http://test") as client: response = await client.put( f"/api/admin/billing/pricing/{config_id}", headers={"Authorization": f"Bearer {admin_token}"}, json={ "config_key": "storage_per_gb_month", "config_value": 0.005 } ) assert response.status_code == status.HTTP_200_OK updated = await PricingConfig.get(id=config_id) assert updated.config_value == Decimal("0.005") @pytest.mark.asyncio async def test_admin_get_stats(admin_user, admin_token): async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( "/api/admin/billing/stats", headers={"Authorization": f"Bearer {admin_token}"} ) assert response.status_code == status.HTTP_200_OK data = response.json() assert "total_revenue" in data assert "total_invoices" in data assert "pending_invoices" in data @pytest.mark.asyncio async def test_non_admin_cannot_access_admin_endpoints(test_user, auth_token): async with AsyncClient(app=app, base_url="http://test") as client: response = await client.get( "/api/admin/billing/pricing", headers={"Authorization": f"Bearer {auth_token}"} ) assert response.status_code == status.HTTP_403_FORBIDDEN