django-simplekeys/simplekeys/tests/test_usagereport.py
2024-12-16 23:04:51 -06:00

48 lines
1.7 KiB
Python

import csv
from django.test import TestCase, override_settings
from django.test.utils import captured_stdout
from django.core.management import call_command
from freezegun import freeze_time
from ..models import Tier, Key, Zone
from ..backends import CacheBackend
class UsageReportTestCase(TestCase):
def setUp(self):
b = CacheBackend()
Zone.objects.create(slug='default', name='Default')
Zone.objects.create(slug='special', name='Special')
tier = Tier.objects.create(slug='default')
Key.objects.create(key='key1', tier=tier, email='key1@example.com')
Key.objects.create(key='key2', tier=tier, email='key2@example.com')
# key, zone, day, num
self.usage = (
# Key 1 usage
('key1', 'default', '20170501', '20'),
('key1', 'default', '20170502', '200'),
('key1', 'default', '20170503', '20'),
('key1', 'special', '20170502', '5'),
# Key 2 usage
('key2', 'special', '20170501', '1'),
('key2', 'special', '20170502', '1'),
('key2', 'special', '20170503', '1'),
)
for key, zone, day, num in self.usage:
for _ in range(int(num)):
b.get_and_inc_quota_value(key, zone, day)
@override_settings(SIMPLEKEYS_RATE_LIMIT_BACKEND='simplekeys.backends.CacheBackend')
def test_basic_report(self):
with freeze_time('2017-05-05'):
with captured_stdout() as stdout:
call_command('usagereport')
stdout.seek(0)
data = tuple(tuple(row) for row in csv.reader(stdout))[1:]
self.assertEquals(len(data), 7)
self.assertEquals(set(data), set(self.usage))