diff --git a/logistics_chat/hooks.py b/logistics_chat/hooks.py index 1530daf..817be6f 100644 --- a/logistics_chat/hooks.py +++ b/logistics_chat/hooks.py @@ -7,11 +7,14 @@ from flask import current_app, jsonify, request def rate_limit(): path = request.path + if request.endpoint == 'pages.index' and request.method in ('GET', 'HEAD'): + return None if path in ('/healthz', '/api/visitor/human/poll') or path.startswith('/static') or path.startswith('/api/human/'): return None ip = request.remote_addr or 'unknown' now = time.time() - hits = current_app.extensions['rate_hits'][ip] + bucket = (ip, 'chat') if request.endpoint == 'chat.stream_chat' else ip + hits = current_app.extensions['rate_hits'][bucket] while hits and now - hits[0] > 60: hits.pop(0) if len(hits) >= current_app.config['RATE_LIMIT_PER_MIN']: diff --git a/tests/test_api.py b/tests/test_api.py index 30b05c6..045ec34 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -359,12 +359,44 @@ class ApiTests(unittest.TestCase): def test_health_check_remains_available_when_business_requests_are_rate_limited(self): self.app.config['RATE_LIMIT_PER_MIN'] = 1 - self.assertEqual(self.client.get('/').status_code, 200) + self.assertEqual(self.client.get('/api/get_history').status_code, 200) self.assertEqual(self.client.get('/api/get_history').status_code, 429) response = self.client.get('/healthz') self.assertEqual(response.status_code, 200) self.assertEqual(response.json, {'status': 'ok'}) + def test_homepage_does_not_consume_or_depend_on_business_rate_budget(self): + self.app.config['RATE_LIMIT_PER_MIN'] = 1 + for method in ('get', 'head'): + response = getattr(self.client, method)('/') + self.assertEqual(response.status_code, 200) + self.assertEqual(response.mimetype, 'text/html') + self.assertEqual(self.client.get('/api/get_history').status_code, 200) + self.assertEqual(self.client.get('/api/get_history').status_code, 429) + for method in ('get', 'head'): + response = getattr(self.client, method)('/') + self.assertEqual(response.status_code, 200) + self.assertEqual(response.mimetype, 'text/html') + + def test_chat_and_supporting_requests_have_separate_limited_budgets(self): + self.app.config['RATE_LIMIT_PER_MIN'] = 1 + response = self.client.post('/stream_chat', json=[]) + self.assertEqual(response.status_code, 200) + self.assertEqual(response.mimetype, 'text/event-stream') + self.assertEqual(self.client.post('/stream_chat', json=[]).status_code, 429) + self.assertEqual(self.client.get('/api/get_history').status_code, 200) + self.assertEqual(self.client.get('/api/get_history').status_code, 429) + self.assertEqual(self.client.get('/stream_chat').status_code, 429) + + def test_supporting_request_limit_does_not_block_chat(self): + self.app.config['RATE_LIMIT_PER_MIN'] = 1 + self.assertEqual(self.client.get('/api/get_history').status_code, 200) + self.assertEqual(self.client.get('/api/get_history').status_code, 429) + response = self.client.post('/stream_chat', json=[]) + self.assertEqual(response.status_code, 200) + self.assertEqual(response.mimetype, 'text/event-stream') + self.assertEqual(self.client.post('/stream_chat', json=[]).status_code, 429) + def test_feedback_stats_and_csv(self): cid = self.conversation() with self.app.app_context():