rbox/tests/billing/test_usage_tracker.py

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