rbox/tests/billing/test_invoice_generator.py

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"))