findEthics/Atlas
0
1#!/usr/bin/env python32"""3Integration tests for chat requests with user authentication4 5This test file focuses on end-to-end testing of the chat API with user_id support,6including request validation, response handling, and data persistence.7"""8 9import sys10import os11sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))12 13# Load environment variables14try:15 from dotenv import load_dotenv16 load_dotenv()17except ImportError:18 pass # dotenv not available, continue without it19 20import asyncio21import httpx22import json23import time24from datetime import datetime25from typing import Optional, Dict, Any26 27 28class TestChatRequestValidation:29 """Test chat request validation with user_id"""30 31 async def test_valid_user_id_formats(self):32 """Test chat requests with various valid user_id formats"""33 valid_user_ids = [34 "user123",35 "user_123",36 "user-123",37 "user_123-test",38 "123user",39 "a", # Single character40 "a" * 255, # Maximum length41 ]42 43 for user_id in valid_user_ids:44 chat_data = {45 "prompt": f"Test message for user {user_id}",46 "max_new_tokens": 50,47 "use_search": False,48 "temperature": 0.7,49 "user_id": user_id50 }51 52 try:53 async with httpx.AsyncClient(timeout=30.0) as client:54 response = await client.post(55 "http://localhost:7860/chat",56 json=chat_data,57 headers={"Content-Type": "application/json"}58 )59 60 assert response.status_code == 200, f"Failed for user_id: {user_id}"61 result = response.json()62 assert "response" in result63 64 # Check session ID in headers65 session_id = response.headers.get('X-Session-ID')66 assert session_id is not None67 68 print(f"✅ Valid user_id '{user_id}' accepted")69 70 except httpx.ConnectError:71 print("⚠️ Server not running - skipping integration test")72 return73 74 async def test_invalid_user_id_formats(self):75 """Test chat requests with invalid user_id formats"""76 invalid_user_ids = [77 "user@123", # @ symbol78 "user 123", # space79 "user.123", # period80 "user#123", # hash81 "user$123", # dollar sign82 "user%123", # percent83 "user&123", # ampersand84 "user*123", # asterisk85 "user+123", # plus86 "user=123", # equals87 "user[123]", # brackets88 "user{123}", # braces89 "user|123", # pipe90 "user\\123", # backslash91 "user/123", # forward slash92 "user:123", # colon93 "user;123", # semicolon94 "user<123>", # angle brackets95 "user?123", # question mark96 "user,123", # comma97 "user'123", # single quote98 'user"123', # double quote99 "user`123", # backtick100 "user~123", # tilde101 "user!123", # exclamation102 "a" * 256, # Too long103 ]104 105 for user_id in invalid_user_ids:106 chat_data = {107 "prompt": f"Test message for invalid user {user_id}",108 "max_new_tokens": 50,109 "use_search": False,110 "temperature": 0.7,111 "user_id": user_id112 }113 114 try:115 async with httpx.AsyncClient(timeout=30.0) as client:116 response = await client.post(117 "http://localhost:7860/chat",118 json=chat_data,119 headers={"Content-Type": "application/json"}120 )121 122 assert response.status_code == 400, f"Should have failed for user_id: {user_id}"123 result = response.json()124 assert "detail" in result125 126 print(f"✅ Invalid user_id '{user_id}' correctly rejected")127 128 except httpx.ConnectError:129 print("⚠️ Server not running - skipping integration test")130 return131 132 async def test_empty_user_id_handling(self):133 """Test that empty user_id is treated as anonymous"""134 empty_user_ids = ["", " ", "\t", "\n"]135 136 for empty_user_id in empty_user_ids:137 chat_data = {138 "prompt": "Test message with empty user_id",139 "max_new_tokens": 50,140 "use_search": False,141 "temperature": 0.7,142 "user_id": empty_user_id143 }144 145 try:146 async with httpx.AsyncClient(timeout=30.0) as client:147 response = await client.post(148 "http://localhost:7860/chat",149 json=chat_data,150 headers={"Content-Type": "application/json"}151 )152 153 assert response.status_code == 200154 result = response.json()155 assert "response" in result156 157 print(f"✅ Empty user_id '{repr(empty_user_id)}' treated as anonymous")158 159 except httpx.ConnectError:160 print("⚠️ Server not running - skipping integration test")161 return162 163 async def test_missing_user_id_field(self):164 """Test that missing user_id field works (backward compatibility)"""165 chat_data = {166 "prompt": "Test message without user_id field",167 "max_new_tokens": 50,168 "use_search": False,169 "temperature": 0.7170 # No user_id field171 }172 173 try:174 async with httpx.AsyncClient(timeout=30.0) as client:175 response = await client.post(176 "http://localhost:7860/chat",177 json=chat_data,178 headers={"Content-Type": "application/json"}179 )180 181 assert response.status_code == 200182 result = response.json()183 assert "response" in result184 185 # Check session ID in headers186 session_id = response.headers.get('X-Session-ID')187 assert session_id is not None188 189 print("✅ Missing user_id field handled correctly")190 191 except httpx.ConnectError:192 print("⚠️ Server not running - skipping integration test")193 return194 195 196class TestChatRequestFlow:197 """Test complete chat request flow with user authentication"""198 199 async def test_authenticated_user_session_flow(self):200 """Test complete flow for authenticated user"""201 user_id = "test_flow_user"202 203 try:204 async with httpx.AsyncClient(timeout=30.0) as client:205 # First request - creates new session206 chat_data1 = {207 "prompt": "First message from authenticated user",208 "max_new_tokens": 50,209 "use_search": False,210 "temperature": 0.7,211 "user_id": user_id212 }213 214 response1 = await client.post(215 "http://localhost:7860/chat",216 json=chat_data1,217 headers={"Content-Type": "application/json"}218 )219 220 assert response1.status_code == 200221 result1 = response1.json()222 assert "response" in result1223 224 session_id = response1.headers.get('X-Session-ID')225 assert session_id is not None226 227 print(f"✅ First request created session: {session_id}")228 229 # Second request - uses existing session230 chat_data2 = {231 "prompt": "Second message from same user",232 "max_new_tokens": 50,233 "use_search": True, # Enable search this time234 "temperature": 0.7,235 "user_id": user_id236 }237 238 response2 = await client.post(239 "http://localhost:7860/chat",240 json=chat_data2,241 headers={242 "Content-Type": "application/json",243 "X-Session-ID": session_id # Provide session ID244 }245 )246 247 assert response2.status_code == 200248 result2 = response2.json()249 assert "response" in result2250 251 # Should return same session ID252 session_id2 = response2.headers.get('X-Session-ID')253 assert session_id2 == session_id254 255 print(f"✅ Second request used same session: {session_id2}")256 257 # Wait for data to be written258 await asyncio.sleep(2)259 260 # Verify data was stored correctly261 await self._verify_session_data(session_id, user_id, expected_messages=2)262 263 except httpx.ConnectError:264 print("⚠️ Server not running - skipping integration test")265 return266 267 async def test_anonymous_user_session_flow(self):268 """Test complete flow for anonymous user"""269 try:270 async with httpx.AsyncClient(timeout=30.0) as client:271 # First request - anonymous user272 chat_data1 = {273 "prompt": "First message from anonymous user",274 "max_new_tokens": 50,275 "use_search": False,276 "temperature": 0.7277 # No user_id field278 }279 280 response1 = await client.post(281 "http://localhost:7860/chat",282 json=chat_data1,283 headers={"Content-Type": "application/json"}284 )285 286 assert response1.status_code == 200287 result1 = response1.json()288 assert "response" in result1289 290 session_id = response1.headers.get('X-Session-ID')291 assert session_id is not None292 293 print(f"✅ Anonymous request created session: {session_id}")294 295 # Second request - same anonymous user296 chat_data2 = {297 "prompt": "Second message from anonymous user",298 "max_new_tokens": 50,299 "use_search": True,300 "temperature": 0.7301 # No user_id field302 }303 304 response2 = await client.post(305 "http://localhost:7860/chat",306 json=chat_data2,307 headers={308 "Content-Type": "application/json",309 "X-Session-ID": session_id310 }311 )312 313 assert response2.status_code == 200314 result2 = response2.json()315 assert "response" in result2316 317 session_id2 = response2.headers.get('X-Session-ID')318 assert session_id2 == session_id319 320 print(f"✅ Anonymous second request used same session: {session_id2}")321 322 # Wait for data to be written323 await asyncio.sleep(2)324 325 # Verify data was stored correctly (user_id should be None)326 await self._verify_session_data(session_id, None, expected_messages=2)327 328 except httpx.ConnectError:329 print("⚠️ Server not running - skipping integration test")330 return331 332 async def test_mixed_user_sessions(self):333 """Test that different users get different sessions"""334 user_id1 = "test_user_1"335 user_id2 = "test_user_2"336 337 try:338 async with httpx.AsyncClient(timeout=30.0) as client:339 # Request from user 1340 chat_data1 = {341 "prompt": "Message from user 1",342 "max_new_tokens": 50,343 "use_search": False,344 "temperature": 0.7,345 "user_id": user_id1346 }347 348 response1 = await client.post(349 "http://localhost:7860/chat",350 json=chat_data1,351 headers={"Content-Type": "application/json"}352 )353 354 assert response1.status_code == 200355 session_id1 = response1.headers.get('X-Session-ID')356 assert session_id1 is not None357 358 # Request from user 2359 chat_data2 = {360 "prompt": "Message from user 2",361 "max_new_tokens": 50,362 "use_search": False,363 "temperature": 0.7,364 "user_id": user_id2365 }366 367 response2 = await client.post(368 "http://localhost:7860/chat",369 json=chat_data2,370 headers={"Content-Type": "application/json"}371 )372 373 assert response2.status_code == 200374 session_id2 = response2.headers.get('X-Session-ID')375 assert session_id2 is not None376 377 # Sessions should be different378 assert session_id1 != session_id2379 380 print(f"✅ User 1 session: {session_id1}")381 print(f"✅ User 2 session: {session_id2}")382 print("✅ Different users got different sessions")383 384 except httpx.ConnectError:385 print("⚠️ Server not running - skipping integration test")386 return387 388 async def _verify_session_data(self, session_id: str, expected_user_id: Optional[str], expected_messages: int):389 """Verify that session data was stored correctly"""390 try:391 from analytics.database import get_sessions_collection, get_messages_collection392 393 sessions_collection = await get_sessions_collection()394 messages_collection = await get_messages_collection()395 396 if sessions_collection is None or messages_collection is None:397 print("⚠️ Database not available - skipping data verification")398 return399 400 # Check session data401 session_doc = await sessions_collection.find_one({"_id": session_id})402 assert session_doc is not None, f"Session {session_id} not found in database"403 assert session_doc.get("user_id") == expected_user_id, f"Expected user_id {expected_user_id}, got {session_doc.get('user_id')}"404 405 # Check message data406 message_docs = await messages_collection.find({"session_id": session_id}).to_list(None)407 assert len(message_docs) == expected_messages, f"Expected {expected_messages} messages, got {len(message_docs)}"408 409 for message_doc in message_docs:410 assert message_doc.get("user_id") == expected_user_id, f"Message user_id mismatch: expected {expected_user_id}, got {message_doc.get('user_id')}"411 412 print(f"✅ Session data verified: user_id={expected_user_id}, messages={len(message_docs)}")413 414 except Exception as e:415 print(f"⚠️ Could not verify session data: {e}")416 417 418class TestChatRequestPerformance:419 """Test performance of chat requests with user authentication"""420 421 async def test_authenticated_request_performance(self):422 """Test performance of authenticated chat requests"""423 user_id = "perf_test_user"424 425 try:426 async with httpx.AsyncClient(timeout=30.0) as client:427 # Warm up428 chat_data = {429 "prompt": "Warmup message",430 "max_new_tokens": 50,431 "use_search": False,432 "temperature": 0.7,433 "user_id": user_id434 }435 436 await client.post(437 "http://localhost:7860/chat",438 json=chat_data,439 headers={"Content-Type": "application/json"}440 )441 442 # Performance test443 num_requests = 5444 total_time = 0445 446 for i in range(num_requests):447 chat_data = {448 "prompt": f"Performance test message {i}",449 "max_new_tokens": 50,450 "use_search": False,451 "temperature": 0.7,452 "user_id": user_id453 }454 455 start_time = time.time()456 response = await client.post(457 "http://localhost:7860/chat",458 json=chat_data,459 headers={"Content-Type": "application/json"}460 )461 end_time = time.time()462 463 assert response.status_code == 200464 request_time = end_time - start_time465 total_time += request_time466 467 print(f"Request {i+1}: {request_time:.2f}s")468 469 avg_time = total_time / num_requests470 print(f"✅ Average request time: {avg_time:.2f}s")471 472 # Performance assertion (requests should be reasonably fast)473 assert avg_time < 10.0, f"Requests too slow: {avg_time:.2f}s average"474 475 except httpx.ConnectError:476 print("⚠️ Server not running - skipping performance test")477 return478 479 async def test_anonymous_vs_authenticated_performance(self):480 """Compare performance between anonymous and authenticated requests"""481 try:482 async with httpx.AsyncClient(timeout=30.0) as client:483 # Test anonymous requests484 anonymous_times = []485 for i in range(3):486 chat_data = {487 "prompt": f"Anonymous performance test {i}",488 "max_new_tokens": 50,489 "use_search": False,490 "temperature": 0.7491 }492 493 start_time = time.time()494 response = await client.post(495 "http://localhost:7860/chat",496 json=chat_data,497 headers={"Content-Type": "application/json"}498 )499 end_time = time.time()500 501 assert response.status_code == 200502 anonymous_times.append(end_time - start_time)503 504 # Test authenticated requests505 authenticated_times = []506 for i in range(3):507 chat_data = {508 "prompt": f"Authenticated performance test {i}",509 "max_new_tokens": 50,510 "use_search": False,511 "temperature": 0.7,512 "user_id": "perf_auth_user"513 }514 515 start_time = time.time()516 response = await client.post(517 "http://localhost:7860/chat",518 json=chat_data,519 headers={"Content-Type": "application/json"}520 )521 end_time = time.time()522 523 assert response.status_code == 200524 authenticated_times.append(end_time - start_time)525 526 avg_anonymous = sum(anonymous_times) / len(anonymous_times)527 avg_authenticated = sum(authenticated_times) / len(authenticated_times)528 529 print(f"✅ Average anonymous request time: {avg_anonymous:.2f}s")530 print(f"✅ Average authenticated request time: {avg_authenticated:.2f}s")531 532 # Performance should be similar (user authentication shouldn't add significant overhead)533 time_difference = abs(avg_authenticated - avg_anonymous)534 assert time_difference < 2.0, f"Too much performance difference: {time_difference:.2f}s"535 536 except httpx.ConnectError:537 print("⚠️ Server not running - skipping performance comparison")538 return539 540 541async def run_integration_tests():542 """Run all integration tests"""543 print("🔗 Running Chat Integration Tests with User Authentication")544 print("=" * 60)545 546 # Test request validation547 validation_test = TestChatRequestValidation()548 await validation_test.test_valid_user_id_formats()549 await validation_test.test_invalid_user_id_formats()550 await validation_test.test_empty_user_id_handling()551 await validation_test.test_missing_user_id_field()552 print("✅ Request validation tests completed")553 554 # Test request flow555 flow_test = TestChatRequestFlow()556 await flow_test.test_authenticated_user_session_flow()557 await flow_test.test_anonymous_user_session_flow()558 await flow_test.test_mixed_user_sessions()559 print("✅ Request flow tests completed")560 561 # Test performance562 perf_test = TestChatRequestPerformance()563 await perf_test.test_authenticated_request_performance()564 await perf_test.test_anonymous_vs_authenticated_performance()565 print("✅ Performance tests completed")566 567 print("\n🎉 ALL INTEGRATION TESTS COMPLETED!")568 569 570if __name__ == "__main__":571 asyncio.run(run_integration_tests())