tests/test_chat_integration_user_auth.py

#!/usr/bin/env python3
"""
Integration tests for chat requests with user authentication

This test file focuses on end-to-end testing of the chat API with user_id support,
including request validation, response handling, and data persistence.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Load environment variables
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # dotenv not available, continue without it

import asyncio
import httpx
import json
import time
from datetime import datetime
from typing import Optional, Dict, Any


class TestChatRequestValidation:
    """Test chat request validation with user_id"""
    
    async def test_valid_user_id_formats(self):
        """Test chat requests with various valid user_id formats"""
        valid_user_ids = [
            "user123",
            "user_123",
            "user-123",
            "user_123-test",
            "123user",
            "a",  # Single character
            "a" * 255,  # Maximum length
        ]
        
        for user_id in valid_user_ids:
            chat_data = {
                "prompt": f"Test message for user {user_id}",
                "max_new_tokens": 50,
                "use_search": False,
                "temperature": 0.7,
                "user_id": user_id
            }
            
            try:
                async with httpx.AsyncClient(timeout=30.0) as client:
                    response = await client.post(
                        "http://localhost:7860/chat",
                        json=chat_data,
                        headers={"Content-Type": "application/json"}
                    )
                    
                    assert response.status_code == 200, f"Failed for user_id: {user_id}"
                    result = response.json()
                    assert "response" in result
                    
                    # Check session ID in headers
                    session_id = response.headers.get('X-Session-ID')
                    assert session_id is not None
                    
                    print(f"✅ Valid user_id '{user_id}' accepted")
                    
            except httpx.ConnectError:
                print("⚠️  Server not running - skipping integration test")
                return
    
    async def test_invalid_user_id_formats(self):
        """Test chat requests with invalid user_id formats"""
        invalid_user_ids = [
            "user@123",      # @ symbol
            "user 123",      # space
            "user.123",      # period
            "user#123",      # hash
            "user$123",      # dollar sign
            "user%123",      # percent
            "user&123",      # ampersand
            "user*123",      # asterisk
            "user+123",      # plus
            "user=123",      # equals
            "user[123]",     # brackets
            "user{123}",     # braces
            "user|123",      # pipe
            "user\\123",     # backslash
            "user/123",      # forward slash
            "user:123",      # colon
            "user;123",      # semicolon
            "user<123>",     # angle brackets
            "user?123",      # question mark
            "user,123",      # comma
            "user'123",      # single quote
            'user"123',      # double quote
            "user`123",      # backtick
            "user~123",      # tilde
            "user!123",      # exclamation
            "a" * 256,       # Too long
        ]
        
        for user_id in invalid_user_ids:
            chat_data = {
                "prompt": f"Test message for invalid user {user_id}",
                "max_new_tokens": 50,
                "use_search": False,
                "temperature": 0.7,
                "user_id": user_id
            }
            
            try:
                async with httpx.AsyncClient(timeout=30.0) as client:
                    response = await client.post(
                        "http://localhost:7860/chat",
                        json=chat_data,
                        headers={"Content-Type": "application/json"}
                    )
                    
                    assert response.status_code == 400, f"Should have failed for user_id: {user_id}"
                    result = response.json()
                    assert "detail" in result
                    
                    print(f"✅ Invalid user_id '{user_id}' correctly rejected")
                    
            except httpx.ConnectError:
                print("⚠️  Server not running - skipping integration test")
                return
    
    async def test_empty_user_id_handling(self):
        """Test that empty user_id is treated as anonymous"""
        empty_user_ids = ["", "   ", "\t", "\n"]
        
        for empty_user_id in empty_user_ids:
            chat_data = {
                "prompt": "Test message with empty user_id",
                "max_new_tokens": 50,
                "use_search": False,
                "temperature": 0.7,
                "user_id": empty_user_id
            }
            
            try:
                async with httpx.AsyncClient(timeout=30.0) as client:
                    response = await client.post(
                        "http://localhost:7860/chat",
                        json=chat_data,
                        headers={"Content-Type": "application/json"}
                    )
                    
                    assert response.status_code == 200
                    result = response.json()
                    assert "response" in result
                    
                    print(f"✅ Empty user_id '{repr(empty_user_id)}' treated as anonymous")
                    
            except httpx.ConnectError:
                print("⚠️  Server not running - skipping integration test")
                return
    
    async def test_missing_user_id_field(self):
        """Test that missing user_id field works (backward compatibility)"""
        chat_data = {
            "prompt": "Test message without user_id field",
            "max_new_tokens": 50,
            "use_search": False,
            "temperature": 0.7
            # No user_id field
        }
        
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data,
                    headers={"Content-Type": "application/json"}
                )
                
                assert response.status_code == 200
                result = response.json()
                assert "response" in result
                
                # Check session ID in headers
                session_id = response.headers.get('X-Session-ID')
                assert session_id is not None
                
                print("✅ Missing user_id field handled correctly")
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return


class TestChatRequestFlow:
    """Test complete chat request flow with user authentication"""
    
    async def test_authenticated_user_session_flow(self):
        """Test complete flow for authenticated user"""
        user_id = "test_flow_user"
        
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                # First request - creates new session
                chat_data1 = {
                    "prompt": "First message from authenticated user",
                    "max_new_tokens": 50,
                    "use_search": False,
                    "temperature": 0.7,
                    "user_id": user_id
                }
                
                response1 = await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data1,
                    headers={"Content-Type": "application/json"}
                )
                
                assert response1.status_code == 200
                result1 = response1.json()
                assert "response" in result1
                
                session_id = response1.headers.get('X-Session-ID')
                assert session_id is not None
                
                print(f"✅ First request created session: {session_id}")
                
                # Second request - uses existing session
                chat_data2 = {
                    "prompt": "Second message from same user",
                    "max_new_tokens": 50,
                    "use_search": True,  # Enable search this time
                    "temperature": 0.7,
                    "user_id": user_id
                }
                
                response2 = await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data2,
                    headers={
                        "Content-Type": "application/json",
                        "X-Session-ID": session_id  # Provide session ID
                    }
                )
                
                assert response2.status_code == 200
                result2 = response2.json()
                assert "response" in result2
                
                # Should return same session ID
                session_id2 = response2.headers.get('X-Session-ID')
                assert session_id2 == session_id
                
                print(f"✅ Second request used same session: {session_id2}")
                
                # Wait for data to be written
                await asyncio.sleep(2)
                
                # Verify data was stored correctly
                await self._verify_session_data(session_id, user_id, expected_messages=2)
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return
    
    async def test_anonymous_user_session_flow(self):
        """Test complete flow for anonymous user"""
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                # First request - anonymous user
                chat_data1 = {
                    "prompt": "First message from anonymous user",
                    "max_new_tokens": 50,
                    "use_search": False,
                    "temperature": 0.7
                    # No user_id field
                }
                
                response1 = await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data1,
                    headers={"Content-Type": "application/json"}
                )
                
                assert response1.status_code == 200
                result1 = response1.json()
                assert "response" in result1
                
                session_id = response1.headers.get('X-Session-ID')
                assert session_id is not None
                
                print(f"✅ Anonymous request created session: {session_id}")
                
                # Second request - same anonymous user
                chat_data2 = {
                    "prompt": "Second message from anonymous user",
                    "max_new_tokens": 50,
                    "use_search": True,
                    "temperature": 0.7
                    # No user_id field
                }
                
                response2 = await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data2,
                    headers={
                        "Content-Type": "application/json",
                        "X-Session-ID": session_id
                    }
                )
                
                assert response2.status_code == 200
                result2 = response2.json()
                assert "response" in result2
                
                session_id2 = response2.headers.get('X-Session-ID')
                assert session_id2 == session_id
                
                print(f"✅ Anonymous second request used same session: {session_id2}")
                
                # Wait for data to be written
                await asyncio.sleep(2)
                
                # Verify data was stored correctly (user_id should be None)
                await self._verify_session_data(session_id, None, expected_messages=2)
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return
    
    async def test_mixed_user_sessions(self):
        """Test that different users get different sessions"""
        user_id1 = "test_user_1"
        user_id2 = "test_user_2"
        
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                # Request from user 1
                chat_data1 = {
                    "prompt": "Message from user 1",
                    "max_new_tokens": 50,
                    "use_search": False,
                    "temperature": 0.7,
                    "user_id": user_id1
                }
                
                response1 = await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data1,
                    headers={"Content-Type": "application/json"}
                )
                
                assert response1.status_code == 200
                session_id1 = response1.headers.get('X-Session-ID')
                assert session_id1 is not None
                
                # Request from user 2
                chat_data2 = {
                    "prompt": "Message from user 2",
                    "max_new_tokens": 50,
                    "use_search": False,
                    "temperature": 0.7,
                    "user_id": user_id2
                }
                
                response2 = await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data2,
                    headers={"Content-Type": "application/json"}
                )
                
                assert response2.status_code == 200
                session_id2 = response2.headers.get('X-Session-ID')
                assert session_id2 is not None
                
                # Sessions should be different
                assert session_id1 != session_id2
                
                print(f"✅ User 1 session: {session_id1}")
                print(f"✅ User 2 session: {session_id2}")
                print("✅ Different users got different sessions")
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return
    
    async def _verify_session_data(self, session_id: str, expected_user_id: Optional[str], expected_messages: int):
        """Verify that session data was stored correctly"""
        try:
            from analytics.database import get_sessions_collection, get_messages_collection
            
            sessions_collection = await get_sessions_collection()
            messages_collection = await get_messages_collection()
            
            if sessions_collection is None or messages_collection is None:
                print("⚠️  Database not available - skipping data verification")
                return
            
            # Check session data
            session_doc = await sessions_collection.find_one({"_id": session_id})
            assert session_doc is not None, f"Session {session_id} not found in database"
            assert session_doc.get("user_id") == expected_user_id, f"Expected user_id {expected_user_id}, got {session_doc.get('user_id')}"
            
            # Check message data
            message_docs = await messages_collection.find({"session_id": session_id}).to_list(None)
            assert len(message_docs) == expected_messages, f"Expected {expected_messages} messages, got {len(message_docs)}"
            
            for message_doc in message_docs:
                assert message_doc.get("user_id") == expected_user_id, f"Message user_id mismatch: expected {expected_user_id}, got {message_doc.get('user_id')}"
            
            print(f"✅ Session data verified: user_id={expected_user_id}, messages={len(message_docs)}")
            
        except Exception as e:
            print(f"⚠️  Could not verify session data: {e}")


class TestChatRequestPerformance:
    """Test performance of chat requests with user authentication"""
    
    async def test_authenticated_request_performance(self):
        """Test performance of authenticated chat requests"""
        user_id = "perf_test_user"
        
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                # Warm up
                chat_data = {
                    "prompt": "Warmup message",
                    "max_new_tokens": 50,
                    "use_search": False,
                    "temperature": 0.7,
                    "user_id": user_id
                }
                
                await client.post(
                    "http://localhost:7860/chat",
                    json=chat_data,
                    headers={"Content-Type": "application/json"}
                )
                
                # Performance test
                num_requests = 5
                total_time = 0
                
                for i in range(num_requests):
                    chat_data = {
                        "prompt": f"Performance test message {i}",
                        "max_new_tokens": 50,
                        "use_search": False,
                        "temperature": 0.7,
                        "user_id": user_id
                    }
                    
                    start_time = time.time()
                    response = await client.post(
                        "http://localhost:7860/chat",
                        json=chat_data,
                        headers={"Content-Type": "application/json"}
                    )
                    end_time = time.time()
                    
                    assert response.status_code == 200
                    request_time = end_time - start_time
                    total_time += request_time
                    
                    print(f"Request {i+1}: {request_time:.2f}s")
                
                avg_time = total_time / num_requests
                print(f"✅ Average request time: {avg_time:.2f}s")
                
                # Performance assertion (requests should be reasonably fast)
                assert avg_time < 10.0, f"Requests too slow: {avg_time:.2f}s average"
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping performance test")
            return
    
    async def test_anonymous_vs_authenticated_performance(self):
        """Compare performance between anonymous and authenticated requests"""
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                # Test anonymous requests
                anonymous_times = []
                for i in range(3):
                    chat_data = {
                        "prompt": f"Anonymous performance test {i}",
                        "max_new_tokens": 50,
                        "use_search": False,
                        "temperature": 0.7
                    }
                    
                    start_time = time.time()
                    response = await client.post(
                        "http://localhost:7860/chat",
                        json=chat_data,
                        headers={"Content-Type": "application/json"}
                    )
                    end_time = time.time()
                    
                    assert response.status_code == 200
                    anonymous_times.append(end_time - start_time)
                
                # Test authenticated requests
                authenticated_times = []
                for i in range(3):
                    chat_data = {
                        "prompt": f"Authenticated performance test {i}",
                        "max_new_tokens": 50,
                        "use_search": False,
                        "temperature": 0.7,
                        "user_id": "perf_auth_user"
                    }
                    
                    start_time = time.time()
                    response = await client.post(
                        "http://localhost:7860/chat",
                        json=chat_data,
                        headers={"Content-Type": "application/json"}
                    )
                    end_time = time.time()
                    
                    assert response.status_code == 200
                    authenticated_times.append(end_time - start_time)
                
                avg_anonymous = sum(anonymous_times) / len(anonymous_times)
                avg_authenticated = sum(authenticated_times) / len(authenticated_times)
                
                print(f"✅ Average anonymous request time: {avg_anonymous:.2f}s")
                print(f"✅ Average authenticated request time: {avg_authenticated:.2f}s")
                
                # Performance should be similar (user authentication shouldn't add significant overhead)
                time_difference = abs(avg_authenticated - avg_anonymous)
                assert time_difference < 2.0, f"Too much performance difference: {time_difference:.2f}s"
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping performance comparison")
            return


async def run_integration_tests():
    """Run all integration tests"""
    print("🔗 Running Chat Integration Tests with User Authentication")
    print("=" * 60)
    
    # Test request validation
    validation_test = TestChatRequestValidation()
    await validation_test.test_valid_user_id_formats()
    await validation_test.test_invalid_user_id_formats()
    await validation_test.test_empty_user_id_handling()
    await validation_test.test_missing_user_id_field()
    print("✅ Request validation tests completed")
    
    # Test request flow
    flow_test = TestChatRequestFlow()
    await flow_test.test_authenticated_user_session_flow()
    await flow_test.test_anonymous_user_session_flow()
    await flow_test.test_mixed_user_sessions()
    print("✅ Request flow tests completed")
    
    # Test performance
    perf_test = TestChatRequestPerformance()
    await perf_test.test_authenticated_request_performance()
    await perf_test.test_anonymous_vs_authenticated_performance()
    print("✅ Performance tests completed")
    
    print("\n🎉 ALL INTEGRATION TESTS COMPLETED!")


if __name__ == "__main__":
    asyncio.run(run_integration_tests())