rbox/tests/billing/test_invoice_generator.py

import pytest
import pytest_asyncio
from decimal import Decimal
from datetime import date, datetime
from mywebdav.models import User
from mywebdav.billing.models import (
    Invoice,
    InvoiceLineItem,
    PricingConfig,
    UsageAggregate,
    UserSubscription,
)
from mywebdav.billing.invoice_generator import InvoiceGenerator


@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_asyncio.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, pricing_config):
    # Update tax rate
    updated = await PricingConfig.filter(config_key="tax_rate_default").update(
        config_value=Decimal("0.21")
    )
    assert updated == 1  # Should update 1 row

    # Verify the update worked
    tax_config = await PricingConfig.get(config_key="tax_rate_default")
    assert tax_config.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")
    )