import pytest from decimal import Decimal from datetime import date, datetime, timedelta from tortoise.contrib.test import initializer, finalizer from mywebdav.models import User, File, Folder from mywebdav.billing.models import UsageRecord, UsageAggregate from mywebdav.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(["mywebdav.models", "mywebdav.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