feat: add billing module with usage tracking, invoice generation, and stripe integration
This commit is contained in:
@@ -0,0 +1,256 @@
|
||||
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
|
||||
@@ -0,0 +1,195 @@
|
||||
import pytest
|
||||
from decimal import Decimal
|
||||
from datetime import date, datetime
|
||||
from tortoise.contrib.test import initializer, finalizer
|
||||
from rbox.models import User
|
||||
from rbox.billing.models import Invoice, InvoiceLineItem, PricingConfig, UsageAggregate, UserSubscription
|
||||
from rbox.billing.invoice_generator import InvoiceGenerator
|
||||
|
||||
@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
|
||||
)
|
||||
yield user
|
||||
await user.delete()
|
||||
|
||||
@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"
|
||||
))
|
||||
configs.append(await PricingConfig.create(
|
||||
config_key="tax_rate_default",
|
||||
config_value=Decimal("0.0"),
|
||||
description="Default tax rate",
|
||||
unit="percentage"
|
||||
))
|
||||
yield configs
|
||||
for config in configs:
|
||||
await config.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_monthly_invoice_with_usage(test_user, pricing_config):
|
||||
today = date.today()
|
||||
|
||||
await UsageAggregate.create(
|
||||
user=test_user,
|
||||
date=today,
|
||||
storage_bytes_avg=1024 ** 3 * 50,
|
||||
storage_bytes_peak=1024 ** 3 * 55,
|
||||
bandwidth_up_bytes=1024 ** 3 * 10,
|
||||
bandwidth_down_bytes=1024 ** 3 * 20
|
||||
)
|
||||
|
||||
invoice = await InvoiceGenerator.generate_monthly_invoice(test_user, today.year, today.month)
|
||||
|
||||
assert invoice is not None
|
||||
assert invoice.user_id == test_user.id
|
||||
assert invoice.status == "draft"
|
||||
assert invoice.total > 0
|
||||
|
||||
line_items = await invoice.line_items.all()
|
||||
assert len(line_items) > 0
|
||||
|
||||
await invoice.delete()
|
||||
await UsageAggregate.filter(user=test_user).delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_monthly_invoice_below_free_tier(test_user, pricing_config):
|
||||
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 * 5,
|
||||
bandwidth_down_bytes=1024 ** 3 * 10
|
||||
)
|
||||
|
||||
invoice = await InvoiceGenerator.generate_monthly_invoice(test_user, today.year, today.month)
|
||||
|
||||
assert invoice is None
|
||||
|
||||
await UsageAggregate.filter(user=test_user).delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_invoice(test_user, pricing_config):
|
||||
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="draft"
|
||||
)
|
||||
|
||||
finalized = await InvoiceGenerator.finalize_invoice(invoice)
|
||||
|
||||
assert finalized.status == "open"
|
||||
|
||||
await finalized.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_invoice_already_finalized(test_user, pricing_config):
|
||||
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"
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await InvoiceGenerator.finalize_invoice(invoice)
|
||||
|
||||
await invoice.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_invoice_paid(test_user, pricing_config):
|
||||
invoice = await Invoice.create(
|
||||
user=test_user,
|
||||
invoice_number="INV-000003-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"
|
||||
)
|
||||
|
||||
paid = await InvoiceGenerator.mark_invoice_paid(invoice)
|
||||
|
||||
assert paid.status == "paid"
|
||||
assert paid.paid_at is not None
|
||||
|
||||
await paid.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoice_with_tax(test_user):
|
||||
await PricingConfig.filter(config_key="tax_rate_default").update(config_value=Decimal("0.21"))
|
||||
|
||||
today = date.today()
|
||||
|
||||
await UsageAggregate.create(
|
||||
user=test_user,
|
||||
date=today,
|
||||
storage_bytes_avg=1024 ** 3 * 50,
|
||||
storage_bytes_peak=1024 ** 3 * 55,
|
||||
bandwidth_up_bytes=1024 ** 3 * 10,
|
||||
bandwidth_down_bytes=1024 ** 3 * 20
|
||||
)
|
||||
|
||||
invoice = await InvoiceGenerator.generate_monthly_invoice(test_user, today.year, today.month)
|
||||
|
||||
assert invoice is not None
|
||||
assert invoice.tax > 0
|
||||
assert invoice.total == invoice.subtotal + invoice.tax
|
||||
|
||||
await invoice.delete()
|
||||
await UsageAggregate.filter(user=test_user).delete()
|
||||
await PricingConfig.filter(config_key="tax_rate_default").update(config_value=Decimal("0.0"))
|
||||
@@ -0,0 +1,188 @@
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from decimal import Decimal
|
||||
from datetime import date, datetime
|
||||
from tortoise.contrib.test import initializer, finalizer
|
||||
from rbox.models import User
|
||||
from rbox.billing.models import (
|
||||
SubscriptionPlan, UserSubscription, UsageRecord, UsageAggregate,
|
||||
Invoice, InvoiceLineItem, PricingConfig, PaymentMethod, BillingEvent
|
||||
)
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def event_loop():
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop_policy().new_event_loop()
|
||||
yield loop
|
||||
loop.close()
|
||||
|
||||
@pytest_asyncio.fixture(scope="module", autouse=True)
|
||||
async def initialize_tests():
|
||||
initializer(["rbox.models", "rbox.billing.models"], db_url="sqlite://:memory:")
|
||||
yield
|
||||
await finalizer()
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def test_user():
|
||||
user = await User.create(
|
||||
username="testuser",
|
||||
email="test@example.com",
|
||||
hashed_password="hashed_password_here",
|
||||
is_active=True
|
||||
)
|
||||
yield user
|
||||
await user.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscription_plan_creation():
|
||||
plan = await SubscriptionPlan.create(
|
||||
name="starter",
|
||||
display_name="Starter Plan",
|
||||
description="Basic storage plan",
|
||||
storage_gb=100,
|
||||
bandwidth_gb=100,
|
||||
price_monthly=Decimal("5.00"),
|
||||
price_yearly=Decimal("50.00")
|
||||
)
|
||||
|
||||
assert plan.name == "starter"
|
||||
assert plan.storage_gb == 100
|
||||
assert plan.price_monthly == Decimal("5.00")
|
||||
await plan.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_subscription_creation(test_user):
|
||||
subscription = await UserSubscription.create(
|
||||
user=test_user,
|
||||
billing_type="pay_as_you_go",
|
||||
status="active"
|
||||
)
|
||||
|
||||
assert subscription.user_id == test_user.id
|
||||
assert subscription.billing_type == "pay_as_you_go"
|
||||
assert subscription.status == "active"
|
||||
await subscription.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_record_creation(test_user):
|
||||
usage = await UsageRecord.create(
|
||||
user=test_user,
|
||||
record_type="storage",
|
||||
amount_bytes=1024 * 1024 * 100,
|
||||
resource_type="file",
|
||||
resource_id=1,
|
||||
idempotency_key="test_key_123"
|
||||
)
|
||||
|
||||
assert usage.user_id == test_user.id
|
||||
assert usage.record_type == "storage"
|
||||
assert usage.amount_bytes == 1024 * 1024 * 100
|
||||
await usage.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_aggregate_creation(test_user):
|
||||
aggregate = await UsageAggregate.create(
|
||||
user=test_user,
|
||||
date=date.today(),
|
||||
storage_bytes_avg=1024 * 1024 * 500,
|
||||
storage_bytes_peak=1024 * 1024 * 600,
|
||||
bandwidth_up_bytes=1024 * 1024 * 50,
|
||||
bandwidth_down_bytes=1024 * 1024 * 100
|
||||
)
|
||||
|
||||
assert aggregate.user_id == test_user.id
|
||||
assert aggregate.storage_bytes_avg == 1024 * 1024 * 500
|
||||
await aggregate.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoice_creation(test_user):
|
||||
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="draft"
|
||||
)
|
||||
|
||||
assert invoice.user_id == test_user.id
|
||||
assert invoice.invoice_number == "INV-000001-202311"
|
||||
assert invoice.total == Decimal("10.00")
|
||||
await invoice.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoice_line_item_creation(test_user):
|
||||
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="draft"
|
||||
)
|
||||
|
||||
line_item = await InvoiceLineItem.create(
|
||||
invoice=invoice,
|
||||
description="Storage usage",
|
||||
quantity=Decimal("100.000000"),
|
||||
unit_price=Decimal("0.100000"),
|
||||
amount=Decimal("10.0000"),
|
||||
item_type="storage"
|
||||
)
|
||||
|
||||
assert line_item.invoice_id == invoice.id
|
||||
assert line_item.description == "Storage usage"
|
||||
assert line_item.amount == Decimal("10.0000")
|
||||
await line_item.delete()
|
||||
await invoice.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pricing_config_creation(test_user):
|
||||
config = 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",
|
||||
updated_by=test_user
|
||||
)
|
||||
|
||||
assert config.config_key == "storage_per_gb_month"
|
||||
assert config.config_value == Decimal("0.0045")
|
||||
await config.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_payment_method_creation(test_user):
|
||||
payment_method = await PaymentMethod.create(
|
||||
user=test_user,
|
||||
stripe_payment_method_id="pm_test_123",
|
||||
type="card",
|
||||
is_default=True,
|
||||
last4="4242",
|
||||
brand="visa",
|
||||
exp_month=12,
|
||||
exp_year=2025
|
||||
)
|
||||
|
||||
assert payment_method.user_id == test_user.id
|
||||
assert payment_method.last4 == "4242"
|
||||
assert payment_method.is_default is True
|
||||
await payment_method.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_billing_event_creation(test_user):
|
||||
event = await BillingEvent.create(
|
||||
user=test_user,
|
||||
event_type="invoice_created",
|
||||
stripe_event_id="evt_test_123",
|
||||
data={"invoice_id": 1},
|
||||
processed=False
|
||||
)
|
||||
|
||||
assert event.user_id == test_user.id
|
||||
assert event.event_type == "invoice_created"
|
||||
assert event.processed is False
|
||||
await event.delete()
|
||||
@@ -0,0 +1,77 @@
|
||||
import pytest
|
||||
|
||||
def test_billing_module_imports():
|
||||
from rbox.billing import models, stripe_client, usage_tracker, invoice_generator, scheduler
|
||||
assert models is not None
|
||||
assert stripe_client is not None
|
||||
assert usage_tracker is not None
|
||||
assert invoice_generator is not None
|
||||
assert scheduler is not None
|
||||
|
||||
def test_billing_models_exist():
|
||||
from rbox.billing.models import (
|
||||
SubscriptionPlan, UserSubscription, UsageRecord, UsageAggregate,
|
||||
Invoice, InvoiceLineItem, PricingConfig, PaymentMethod, BillingEvent
|
||||
)
|
||||
assert SubscriptionPlan is not None
|
||||
assert UserSubscription is not None
|
||||
assert UsageRecord is not None
|
||||
assert UsageAggregate is not None
|
||||
assert Invoice is not None
|
||||
assert InvoiceLineItem is not None
|
||||
assert PricingConfig is not None
|
||||
assert PaymentMethod is not None
|
||||
assert BillingEvent is not None
|
||||
|
||||
def test_stripe_client_exists():
|
||||
from rbox.billing.stripe_client import StripeClient
|
||||
assert StripeClient is not None
|
||||
assert hasattr(StripeClient, 'create_customer')
|
||||
assert hasattr(StripeClient, 'create_invoice')
|
||||
assert hasattr(StripeClient, 'finalize_invoice')
|
||||
|
||||
def test_usage_tracker_exists():
|
||||
from rbox.billing.usage_tracker import UsageTracker
|
||||
assert UsageTracker is not None
|
||||
assert hasattr(UsageTracker, 'track_storage')
|
||||
assert hasattr(UsageTracker, 'track_bandwidth')
|
||||
assert hasattr(UsageTracker, 'aggregate_daily_usage')
|
||||
assert hasattr(UsageTracker, 'get_current_storage')
|
||||
assert hasattr(UsageTracker, 'get_monthly_usage')
|
||||
|
||||
def test_invoice_generator_exists():
|
||||
from rbox.billing.invoice_generator import InvoiceGenerator
|
||||
assert InvoiceGenerator is not None
|
||||
assert hasattr(InvoiceGenerator, 'generate_monthly_invoice')
|
||||
assert hasattr(InvoiceGenerator, 'finalize_invoice')
|
||||
assert hasattr(InvoiceGenerator, 'mark_invoice_paid')
|
||||
|
||||
def test_scheduler_exists():
|
||||
from rbox.billing.scheduler import scheduler, start_scheduler, stop_scheduler
|
||||
assert scheduler is not None
|
||||
assert callable(start_scheduler)
|
||||
assert callable(stop_scheduler)
|
||||
|
||||
def test_routers_exist():
|
||||
from rbox.routers import billing, admin_billing
|
||||
assert billing is not None
|
||||
assert admin_billing is not None
|
||||
assert hasattr(billing, 'router')
|
||||
assert hasattr(admin_billing, 'router')
|
||||
|
||||
def test_middleware_exists():
|
||||
from rbox.middleware.usage_tracking import UsageTrackingMiddleware
|
||||
assert UsageTrackingMiddleware is not None
|
||||
|
||||
def test_settings_updated():
|
||||
from rbox.settings import settings
|
||||
assert hasattr(settings, 'STRIPE_SECRET_KEY')
|
||||
assert hasattr(settings, 'STRIPE_PUBLISHABLE_KEY')
|
||||
assert hasattr(settings, 'STRIPE_WEBHOOK_SECRET')
|
||||
assert hasattr(settings, 'BILLING_ENABLED')
|
||||
|
||||
def test_main_includes_billing():
|
||||
from rbox.main import app
|
||||
routes = [route.path for route in app.routes]
|
||||
billing_routes = [r for r in routes if '/billing' in r]
|
||||
assert len(billing_routes) > 0
|
||||
@@ -0,0 +1,175 @@
|
||||
import pytest
|
||||
from decimal import Decimal
|
||||
from datetime import date, datetime, timedelta
|
||||
from tortoise.contrib.test import initializer, finalizer
|
||||
from rbox.models import User, File, Folder
|
||||
from rbox.billing.models import UsageRecord, UsageAggregate
|
||||
from rbox.billing.usage_tracker import UsageTracker
|
||||
|
||||
@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
|
||||
)
|
||||
yield user
|
||||
await user.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_track_storage(test_user):
|
||||
await UsageTracker.track_storage(
|
||||
user=test_user,
|
||||
amount_bytes=1024 * 1024 * 100,
|
||||
resource_type="file",
|
||||
resource_id=1
|
||||
)
|
||||
|
||||
records = await UsageRecord.filter(user=test_user, record_type="storage").all()
|
||||
assert len(records) == 1
|
||||
assert records[0].amount_bytes == 1024 * 1024 * 100
|
||||
|
||||
await records[0].delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_track_bandwidth_upload(test_user):
|
||||
await UsageTracker.track_bandwidth(
|
||||
user=test_user,
|
||||
amount_bytes=1024 * 1024 * 50,
|
||||
direction="up",
|
||||
resource_type="file",
|
||||
resource_id=1
|
||||
)
|
||||
|
||||
records = await UsageRecord.filter(user=test_user, record_type="bandwidth_up").all()
|
||||
assert len(records) == 1
|
||||
assert records[0].amount_bytes == 1024 * 1024 * 50
|
||||
|
||||
await records[0].delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_track_bandwidth_download(test_user):
|
||||
await UsageTracker.track_bandwidth(
|
||||
user=test_user,
|
||||
amount_bytes=1024 * 1024 * 75,
|
||||
direction="down",
|
||||
resource_type="file",
|
||||
resource_id=1
|
||||
)
|
||||
|
||||
records = await UsageRecord.filter(user=test_user, record_type="bandwidth_down").all()
|
||||
assert len(records) == 1
|
||||
assert records[0].amount_bytes == 1024 * 1024 * 75
|
||||
|
||||
await records[0].delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregate_daily_usage(test_user):
|
||||
await UsageTracker.track_storage(
|
||||
user=test_user,
|
||||
amount_bytes=1024 * 1024 * 100
|
||||
)
|
||||
|
||||
await UsageTracker.track_bandwidth(
|
||||
user=test_user,
|
||||
amount_bytes=1024 * 1024 * 50,
|
||||
direction="up"
|
||||
)
|
||||
|
||||
await UsageTracker.track_bandwidth(
|
||||
user=test_user,
|
||||
amount_bytes=1024 * 1024 * 75,
|
||||
direction="down"
|
||||
)
|
||||
|
||||
aggregate = await UsageTracker.aggregate_daily_usage(test_user, date.today())
|
||||
|
||||
assert aggregate is not None
|
||||
assert aggregate.storage_bytes_avg > 0
|
||||
assert aggregate.bandwidth_up_bytes == 1024 * 1024 * 50
|
||||
assert aggregate.bandwidth_down_bytes == 1024 * 1024 * 75
|
||||
|
||||
await aggregate.delete()
|
||||
await UsageRecord.filter(user=test_user).delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_storage(test_user):
|
||||
folder = await Folder.create(
|
||||
name="Test Folder",
|
||||
owner=test_user
|
||||
)
|
||||
|
||||
file1 = await File.create(
|
||||
name="test1.txt",
|
||||
path="/storage/test1.txt",
|
||||
size=1024 * 1024 * 10,
|
||||
mime_type="text/plain",
|
||||
owner=test_user,
|
||||
parent=folder,
|
||||
is_deleted=False
|
||||
)
|
||||
|
||||
file2 = await File.create(
|
||||
name="test2.txt",
|
||||
path="/storage/test2.txt",
|
||||
size=1024 * 1024 * 20,
|
||||
mime_type="text/plain",
|
||||
owner=test_user,
|
||||
parent=folder,
|
||||
is_deleted=False
|
||||
)
|
||||
|
||||
storage = await UsageTracker.get_current_storage(test_user)
|
||||
|
||||
assert storage == 1024 * 1024 * 30
|
||||
|
||||
await file1.delete()
|
||||
await file2.delete()
|
||||
await folder.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_monthly_usage(test_user):
|
||||
today = date.today()
|
||||
|
||||
aggregate = await UsageAggregate.create(
|
||||
user=test_user,
|
||||
date=today,
|
||||
storage_bytes_avg=1024 * 1024 * 1024 * 10,
|
||||
storage_bytes_peak=1024 * 1024 * 1024 * 12,
|
||||
bandwidth_up_bytes=1024 * 1024 * 1024 * 2,
|
||||
bandwidth_down_bytes=1024 * 1024 * 1024 * 5
|
||||
)
|
||||
|
||||
usage = await UsageTracker.get_monthly_usage(test_user, today.year, today.month)
|
||||
|
||||
assert usage["storage_gb_avg"] == pytest.approx(10.0, rel=0.01)
|
||||
assert usage["storage_gb_peak"] == pytest.approx(12.0, rel=0.01)
|
||||
assert usage["bandwidth_up_gb"] == pytest.approx(2.0, rel=0.01)
|
||||
assert usage["bandwidth_down_gb"] == pytest.approx(5.0, rel=0.01)
|
||||
|
||||
await aggregate.delete()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_monthly_usage_empty(test_user):
|
||||
future_date = date.today() + timedelta(days=365)
|
||||
|
||||
usage = await UsageTracker.get_monthly_usage(test_user, future_date.year, future_date.month)
|
||||
|
||||
assert usage["storage_gb_avg"] == 0
|
||||
assert usage["storage_gb_peak"] == 0
|
||||
assert usage["bandwidth_up_gb"] == 0
|
||||
assert usage["bandwidth_down_gb"] == 0
|
||||
Reference in New Issue
Block a user