| import concurrent.futures |
| import json |
| import tempfile |
| import threading |
| import unittest |
| import urllib.request |
| import urllib.error |
| from pathlib import Path |
| from http.server import ThreadingHTTPServer |
| from unittest.mock import patch |
| from api_access import AccessStore, AccessError |
| import serve_models |
|
|
|
|
| class AccessTests(unittest.TestCase): |
| def setUp(self): |
| self.temp = tempfile.TemporaryDirectory() |
| self.store = AccessStore(Path(self.temp.name)/'access.db') |
|
|
| def tearDown(self): |
| self.temp.cleanup() |
|
|
| def test_approval_revoke_and_hashed_storage(self): |
| r = self.store.request('Test app', 'owner@example.test') |
| with self.assertRaises(AccessError) as e: self.store.authorize(r['api_key']) |
| self.assertEqual(e.exception.status, 403) |
| self.store.change(r['request_id'], 'approved', 2, 3) |
| self.store.authorize(r['api_key']) |
| self.store.change(r['request_id'], 'revoked') |
| with self.assertRaises(AccessError) as e: self.store.authorize(r['api_key']) |
| self.assertEqual(e.exception.status, 403) |
| self.assertNotIn(r['api_key'].encode(), self.store.path.read_bytes()) |
| self.assertEqual(len(AccessStore(self.store.path).list()), 1) |
|
|
| def test_atomic_rate_limit(self): |
| def charge(_): |
| try: self.store.charge('shared', 2, 10, now=200000); return True |
| except AccessError: return False |
| with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: |
| self.assertEqual(sum(executor.map(charge, range(20))), 2) |
| self.store.charge('shared', 2, 10, now=200061) |
|
|
| def test_daily_reset(self): |
| self.store.charge('daily', 10, 1, now=200000) |
| with self.assertRaises(AccessError): self.store.charge('daily', 10, 1, now=200100) |
| self.store.charge('daily', 10, 1, now=259201) |
|
|
| def test_http_pending_approved_rate_limit_and_revocation(self): |
| serve_models.ACCESS = self.store |
| serve_models.ALLOW_ANONYMOUS = False |
| server = ThreadingHTTPServer(('127.0.0.1', 0), serve_models.Handler) |
| thread = threading.Thread(target=server.serve_forever, daemon=True); thread.start() |
| def post(path, data, key=None): |
| headers = {'Content-Type':'application/json'} |
| if key: headers['Authorization'] = 'Bearer '+key |
| req = urllib.request.Request(f'http://127.0.0.1:{server.server_port}'+path, |
| data=json.dumps(data).encode(),headers=headers) |
| try: response = urllib.request.urlopen(req) |
| except urllib.error.HTTPError as error: response = error |
| with response: return response.status, json.load(response), response.headers |
| try: |
| self.assertEqual(post('/v1/emotions', {'text':'Hello'})[0],401) |
| status,r,_ = post('/v1/access-requests',{'app_name':'HTTP test','contact':'test@example.test'}) |
| self.assertEqual(status,201) |
| self.assertEqual(post('/v1/emotions',{'text':'Hello'},r['api_key'])[0],403) |
| self.store.change(r['request_id'],'approved',2,2) |
| with patch.object(serve_models,'predict',return_value={'windows':[]}): |
| for _ in range(2): self.assertEqual(post('/v1/emotions',{'text':'Hello'},r['api_key'])[0],200) |
| status,_,headers = post('/v1/emotions',{'text':'Hello'},r['api_key']) |
| self.assertEqual(status,429); self.assertGreater(int(headers['Retry-After']),0) |
| self.store.change(r['request_id'],'revoked') |
| self.assertEqual(post('/v1/emotions',{'text':'Hello'},r['api_key'])[0],403) |
| finally: |
| server.shutdown(); server.server_close(); thread.join() |
|
|
|
|
| if __name__ == '__main__': unittest.main() |
|
|
