tests/test_user_authentication_comprehensive.py

#!/usr/bin/env python3
"""
Comprehensive tests for user authentication feature

This test suite covers:
1. Unit tests for updated models with user_id validation
2. Integration tests for chat requests with and without user_id
3. Tests for user-specific analytics functions
4. Backward compatibility tests for anonymous users
5. Performance tests for user_id queries and indexes
"""

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 time
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
import uuid

# Import models and functions to test
from analytics.models import Session, Message, SearchAnalytics
from analytics.collectors import create_session, track_message, track_search
from analytics.dashboard import (
    get_user_statistics,
    get_user_analytics,
    get_authenticated_vs_anonymous_metrics,
    get_basic_stats
)
from analytics.database import (
    get_sessions_collection,
    get_messages_collection,
    get_search_analytics_collection
)


class TestUserIdValidation:
    """Unit tests for user_id validation in models"""
    
    def test_session_valid_user_id(self):
        """Test Session model with valid user_id values"""
        # Valid user_id
        session = Session(user_id="user_123")
        assert session.user_id == "user_123"
        
        # Valid user_id with hyphens and underscores
        session = Session(user_id="user-123_test")
        assert session.user_id == "user-123_test"
        
        # None user_id (anonymous)
        session = Session(user_id=None)
        assert session.user_id is None
        
        # Empty string becomes None
        session = Session(user_id="")
        assert session.user_id is None
        
        # Whitespace-only string becomes None
        session = Session(user_id="   ")
        assert session.user_id is None
    
    def test_session_invalid_user_id(self):
        """Test Session model with invalid user_id values"""
        # Non-string user_id
        try:
            Session(user_id=123)
            assert False, "Should have raised error for non-string user_id"
        except Exception as e:
            assert "string" in str(e).lower()
        
        # Too long user_id
        try:
            Session(user_id="a" * 256)
            assert False, "Should have raised error for too long user_id"
        except Exception as e:
            assert "255" in str(e)
        
        # Invalid characters
        for invalid_id in ["user@123", "user 123", "user.123"]:
            try:
                Session(user_id=invalid_id)
                assert False, f"Should have raised error for {invalid_id}"
            except Exception as e:
                assert "alphanumeric" in str(e).lower()
    
    def test_message_valid_user_id(self):
        """Test Message model with valid user_id values"""
        # Valid user_id
        message = Message(
            session_id="test_session",
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            user_id="user_123"
        )
        assert message.user_id == "user_123"
        
        # None user_id (anonymous)
        message = Message(
            session_id="test_session",
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            user_id=None
        )
        assert message.user_id is None
    
    def test_message_invalid_user_id(self):
        """Test Message model with invalid user_id values"""
        # Non-string user_id
        try:
            Message(
                session_id="test_session",
                prompt_length=50,
                response_length=100,
                response_time_ms=1000,
                user_id=123
            )
            assert False, "Should have raised error for non-string user_id"
        except Exception as e:
            assert "string" in str(e).lower()
    
    def test_search_analytics_valid_user_id(self):
        """Test SearchAnalytics model with valid user_id values"""
        # Valid user_id
        search = SearchAnalytics(
            message_id="test_message",
            search_query="test query",
            user_id="user_123"
        )
        assert search.user_id == "user_123"
        
        # None user_id (anonymous)
        search = SearchAnalytics(
            message_id="test_message",
            search_query="test query",
            user_id=None
        )
        assert search.user_id is None
    
    def test_search_analytics_invalid_user_id(self):
        """Test SearchAnalytics model with invalid user_id values"""
        # Non-string user_id
        try:
            SearchAnalytics(
                message_id="test_message",
                search_query="test query",
                user_id=123
            )
            assert False, "Should have raised error for non-string user_id"
        except Exception as e:
            assert "string" in str(e).lower()
    
    def test_model_to_dict_includes_user_id(self):
        """Test that to_dict() methods include user_id field"""
        # Session with user_id
        session = Session(user_id="user_123")
        session_dict = session.to_dict()
        assert "user_id" in session_dict
        assert session_dict["user_id"] == "user_123"
        
        # Message with user_id
        message = Message(
            session_id="test_session",
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            user_id="user_123"
        )
        message_dict = message.to_dict()
        assert "user_id" in message_dict
        assert message_dict["user_id"] == "user_123"
        
        # SearchAnalytics with user_id
        search = SearchAnalytics(
            message_id="test_message",
            search_query="test query",
            user_id="user_123"
        )
        search_dict = search.to_dict()
        assert "user_id" in search_dict
        assert search_dict["user_id"] == "user_123"


class TestAnalyticsCollectors:
    """Unit tests for analytics collectors with user_id support"""
    
    async def test_create_session_with_user_id(self):
        """Test create_session function with user_id"""
        # Create session with user_id
        session = await create_session(user_agent="TestAgent", user_id="user_123")
        assert session.user_id == "user_123"
        assert session.user_agent == "TestAgent"
        
        # Create anonymous session
        session = await create_session(user_agent="TestAgent", user_id=None)
        assert session.user_id is None
        
        # Create session without user_id parameter
        session = await create_session(user_agent="TestAgent")
        assert session.user_id is None
    
    async def test_track_message_with_user_id(self):
        """Test track_message function with user_id"""
        # Create a session first
        session = await create_session(user_id="user_123")
        
        # Track message with user_id
        message = await track_message(
            session_id=session.session_id,
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            user_id="user_123"
        )
        
        assert message is not None
        assert message.user_id == "user_123"
        assert message.session_id == session.session_id
    
    async def test_track_message_user_id_mismatch_warning(self):
        """Test that user_id mismatch between session and message logs warning"""
        # Create a session with one user_id
        session = await create_session(user_id="user_123")
        
        # Track message with different user_id
        message = await track_message(
            session_id=session.session_id,
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            user_id="user_456"  # Different user_id
        )
        
        assert message is not None
        assert message.user_id == "user_456"  # Message should use provided user_id
        
        # Note: In a real test environment, we would check logs
        # For now, we just verify the message was created with the provided user_id
        print("✅ User ID mismatch handling tested (warning would be logged)")
    
    async def test_track_search_with_user_id(self):
        """Test track_search function with user_id"""
        # Create session and message first
        session = await create_session(user_id="user_123")
        message = await track_message(
            session_id=session.session_id,
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            user_id="user_123"
        )
        
        # Track search with user_id
        search = await track_search(
            message_id=message.message_id,
            search_query="test query",
            search_terms=["test", "query"],
            brave_results=5,
            duckduckgo_results=3,
            total_unique_results=7,
            user_id="user_123"
        )
        
        assert search is not None
        assert search.user_id == "user_123"
        assert search.message_id == message.message_id


class TestChatIntegration:
    """Integration tests for chat requests with user_id"""
    
    async def test_chat_request_with_user_id(self):
        """Test chat request with user_id parameter"""
        chat_data = {
            "prompt": "Test message with user authentication",
            "max_new_tokens": 100,
            "use_search": False,
            "temperature": 0.7,
            "user_id": "test_user_123"
        }
        
        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
                
                return session_id
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return None
    
    async def test_chat_request_without_user_id(self):
        """Test chat request without user_id parameter (anonymous)"""
        chat_data = {
            "prompt": "Test anonymous message",
            "max_new_tokens": 100,
            "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
                
                return session_id
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return None
    
    async def test_chat_request_invalid_user_id(self):
        """Test chat request with invalid user_id"""
        chat_data = {
            "prompt": "Test message with invalid user_id",
            "max_new_tokens": 100,
            "use_search": False,
            "temperature": 0.7,
            "user_id": "invalid@user"  # Invalid characters
        }
        
        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
                result = response.json()
                assert "detail" in result
                assert "user_id can only contain alphanumeric characters" in result["detail"]
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return
    
    async def test_chat_request_empty_user_id(self):
        """Test chat request with empty user_id (should be treated as anonymous)"""
        chat_data = {
            "prompt": "Test message with empty user_id",
            "max_new_tokens": 100,
            "use_search": False,
            "temperature": 0.7,
            "user_id": ""  # Empty string
        }
        
        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
                
        except httpx.ConnectError:
            print("⚠️  Server not running - skipping integration test")
            return


class TestUserAnalyticsFunctions:
    """Tests for user-specific analytics functions"""
    
    async def setup_test_data(self):
        """Set up test data for analytics tests"""
        # Create test sessions and messages
        auth_session = await create_session(user_id="test_user_analytics")
        anon_session = await create_session(user_id=None)
        
        # Track messages
        await track_message(
            session_id=auth_session.session_id,
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            used_search=True,
            user_id="test_user_analytics"
        )
        
        await track_message(
            session_id=anon_session.session_id,
            prompt_length=40,
            response_length=80,
            response_time_ms=800,
            used_search=False,
            user_id=None
        )
        
        return {
            "auth_session": auth_session,
            "anon_session": anon_session
        }
    

    async def test_get_user_statistics(self):
        """Test get_user_statistics function"""
        result = await get_user_statistics()
        
        assert isinstance(result, dict)
        required_keys = [
            "total_sessions", "authenticated_sessions", "anonymous_sessions",
            "authenticated_session_percentage", "total_messages", 
            "authenticated_messages", "anonymous_messages",
            "authenticated_message_percentage", "unique_authenticated_users"
        ]
        
        for key in required_keys:
            assert key in result
            assert isinstance(result[key], (int, float))
            assert result[key] >= 0
    

    async def test_get_user_analytics_valid_user(self):
        """Test get_user_analytics with valid user_id"""
        result = await get_user_analytics("test_user_analytics")
        
        assert isinstance(result, dict)
        assert result.get("user_id") == "test_user_analytics"
        
        required_keys = [
            "total_sessions", "active_sessions", "total_messages",
            "messages_with_search", "search_usage_percentage",
            "avg_response_time_ms", "daily_activity_last_30_days"
        ]
        
        for key in required_keys:
            assert key in result
    

    async def test_get_user_analytics_invalid_user(self):
        """Test get_user_analytics with invalid user_id"""
        # Test with None
        result = await get_user_analytics(None)
        assert "error" in result
        
        # Test with empty string
        result = await get_user_analytics("")
        assert "error" in result
    

    async def test_get_authenticated_vs_anonymous_metrics(self):
        """Test get_authenticated_vs_anonymous_metrics function"""
        result = await get_authenticated_vs_anonymous_metrics()
        
        assert isinstance(result, dict)
        assert "authenticated" in result
        assert "anonymous" in result
        assert "comparison" in result
        
        # Check structure of authenticated metrics
        auth_metrics = result["authenticated"]
        assert "sessions" in auth_metrics
        assert "messages" in auth_metrics
        assert "avg_messages_per_session" in auth_metrics
        
        # Check structure of anonymous metrics
        anon_metrics = result["anonymous"]
        assert "sessions" in anon_metrics
        assert "messages" in anon_metrics
        assert "avg_messages_per_session" in anon_metrics
    

    async def test_basic_stats_with_user_filter(self):
        """Test get_basic_stats with user_id filter"""
        # Test without filter
        result_all = await get_basic_stats()
        assert isinstance(result_all, dict)
        
        # Test with user filter
        result_filtered = await get_basic_stats(user_id="test_user_analytics")
        assert isinstance(result_filtered, dict)
        assert "filtered_by_user_id" in result_filtered
        assert result_filtered["filtered_by_user_id"] == "test_user_analytics"


class TestBackwardCompatibility:
    """Tests for backward compatibility with anonymous users"""
    

    async def test_anonymous_session_creation(self):
        """Test that anonymous sessions work as before"""
        # Create session without user_id (old way)
        session = await create_session(user_agent="TestAgent")
        assert session.user_id is None
        assert session.user_agent == "TestAgent"
        
        # Create session with explicit None user_id
        session = await create_session(user_agent="TestAgent", user_id=None)
        assert session.user_id is None
    

    async def test_anonymous_message_tracking(self):
        """Test that anonymous message tracking works as before"""
        session = await create_session()
        
        # Track message without user_id (old way)
        message = await track_message(
            session_id=session.session_id,
            prompt_length=50,
            response_length=100,
            response_time_ms=1000
        )
        
        assert message is not None
        assert message.user_id is None
        assert message.session_id == session.session_id
    

    async def test_anonymous_search_tracking(self):
        """Test that anonymous search tracking works as before"""
        session = await create_session()
        message = await track_message(
            session_id=session.session_id,
            prompt_length=50,
            response_length=100,
            response_time_ms=1000
        )
        
        # Track search without user_id (old way)
        search = await track_search(
            message_id=message.message_id,
            search_query="test query",
            search_terms=["test", "query"]
        )
        
        assert search is not None
        assert search.user_id is None
    

    async def test_existing_analytics_functions_work(self):
        """Test that existing analytics functions work with mixed data"""
        # Create both authenticated and anonymous data
        auth_session = await create_session(user_id="test_user")
        anon_session = await create_session()
        
        await track_message(
            session_id=auth_session.session_id,
            prompt_length=50,
            response_length=100,
            response_time_ms=1000,
            user_id="test_user"
        )
        
        await track_message(
            session_id=anon_session.session_id,
            prompt_length=40,
            response_length=80,
            response_time_ms=800
        )
        
        # Test that basic stats work
        stats = await get_basic_stats()
        assert isinstance(stats, dict)
        assert stats["total_sessions"] >= 2
        assert stats["total_messages"] >= 2


class TestPerformance:
    """Performance tests for user_id queries and indexes"""
    

    async def test_user_id_query_performance(self):
        """Test performance of user_id queries"""
        # Create test data
        user_ids = [f"perf_user_{i}" for i in range(10)]
        sessions = []
        
        # Create sessions for performance testing
        start_time = time.time()
        for user_id in user_ids:
            session = await create_session(user_id=user_id)
            sessions.append(session)
            
            # Track multiple messages per session
            for j in range(5):
                await track_message(
                    session_id=session.session_id,
                    prompt_length=50,
                    response_length=100,
                    response_time_ms=1000,
                    user_id=user_id
                )
        
        creation_time = time.time() - start_time
        print(f"Data creation took: {creation_time:.2f} seconds")
        
        # Test query performance
        sessions_collection = await get_sessions_collection()
        messages_collection = await get_messages_collection()
        
        if sessions_collection and messages_collection:
            # Test user-specific session queries
            start_time = time.time()
            for user_id in user_ids:
                user_sessions = await sessions_collection.count_documents({"user_id": user_id})
                assert user_sessions == 1
            
            session_query_time = time.time() - start_time
            print(f"Session queries took: {session_query_time:.2f} seconds")
            
            # Test user-specific message queries
            start_time = time.time()
            for user_id in user_ids:
                user_messages = await messages_collection.count_documents({"user_id": user_id})
                assert user_messages == 5
            
            message_query_time = time.time() - start_time
            print(f"Message queries took: {message_query_time:.2f} seconds")
            
            # Performance assertions (queries should be reasonably fast)
            assert session_query_time < 5.0, f"Session queries too slow: {session_query_time:.2f}s"
            assert message_query_time < 5.0, f"Message queries too slow: {message_query_time:.2f}s"
    

    async def test_compound_index_performance(self):
        """Test performance of compound (user_id, timestamp) queries"""
        # Create test data with timestamps
        user_id = "compound_test_user"
        session = await create_session(user_id=user_id)
        
        # Create messages over time
        start_time = time.time()
        for i in range(20):
            await track_message(
                session_id=session.session_id,
                prompt_length=50,
                response_length=100,
                response_time_ms=1000,
                user_id=user_id
            )
            # Small delay to create different timestamps
            await asyncio.sleep(0.01)
        
        creation_time = time.time() - start_time
        print(f"Compound test data creation took: {creation_time:.2f} seconds")
        
        # Test compound queries (user_id + timestamp range)
        messages_collection = await get_messages_collection()
        if messages_collection:
            now = datetime.utcnow()
            one_hour_ago = now - timedelta(hours=1)
            
            start_time = time.time()
            recent_messages = await messages_collection.count_documents({
                "user_id": user_id,
                "timestamp": {"$gte": one_hour_ago}
            })
            
            compound_query_time = time.time() - start_time
            print(f"Compound query took: {compound_query_time:.2f} seconds")
            
            assert recent_messages == 20
            assert compound_query_time < 2.0, f"Compound query too slow: {compound_query_time:.2f}s"
    

    async def test_analytics_function_performance(self):
        """Test performance of user analytics functions"""
        # Create test user with data
        user_id = "analytics_perf_user"
        session = await create_session(user_id=user_id)
        
        # Create multiple messages
        for i in range(50):
            await track_message(
                session_id=session.session_id,
                prompt_length=50,
                response_length=100,
                response_time_ms=1000,
                used_search=(i % 2 == 0),  # Alternate search usage
                user_id=user_id
            )
        
        # Test performance of user analytics function
        start_time = time.time()
        user_analytics = await get_user_analytics(user_id)
        analytics_time = time.time() - start_time
        
        print(f"User analytics query took: {analytics_time:.2f} seconds")
        
        assert isinstance(user_analytics, dict)
        assert user_analytics.get("user_id") == user_id
        assert user_analytics.get("total_messages") == 50
        assert analytics_time < 5.0, f"User analytics too slow: {analytics_time:.2f}s"


# Test runner functions
async def run_unit_tests():
    """Run unit tests"""
    print("🧪 Running Unit Tests")
    print("=" * 50)
    
    # Test user_id validation
    test_validation = TestUserIdValidation()
    test_validation.test_session_valid_user_id()
    test_validation.test_session_invalid_user_id()
    test_validation.test_message_valid_user_id()
    test_validation.test_message_invalid_user_id()
    test_validation.test_search_analytics_valid_user_id()
    test_validation.test_search_analytics_invalid_user_id()
    test_validation.test_model_to_dict_includes_user_id()
    print("✅ User ID validation tests passed")
    
    # Test analytics collectors
    test_collectors = TestAnalyticsCollectors()
    await test_collectors.test_create_session_with_user_id()
    await test_collectors.test_track_message_with_user_id()
    await test_collectors.test_track_search_with_user_id()
    print("✅ Analytics collectors tests passed")


async def run_integration_tests():
    """Run integration tests"""
    print("\n🔗 Running Integration Tests")
    print("=" * 50)
    
    test_integration = TestChatIntegration()
    try:
        await test_integration.test_chat_request_with_user_id()
        await test_integration.test_chat_request_without_user_id()
        await test_integration.test_chat_request_invalid_user_id()
        await test_integration.test_chat_request_empty_user_id()
        print("✅ Chat integration tests passed")
    except Exception as e:
        print(f"⚠️  Integration tests skipped: {e}")


async def run_analytics_tests():
    """Run analytics function tests"""
    print("\n📊 Running Analytics Function Tests")
    print("=" * 50)
    
    test_analytics = TestUserAnalyticsFunctions()
    await test_analytics.test_get_user_statistics()
    await test_analytics.test_get_user_analytics_valid_user()
    await test_analytics.test_get_user_analytics_invalid_user()
    await test_analytics.test_get_authenticated_vs_anonymous_metrics()
    await test_analytics.test_basic_stats_with_user_filter()
    print("✅ Analytics function tests passed")


async def run_compatibility_tests():
    """Run backward compatibility tests"""
    print("\n🔄 Running Backward Compatibility Tests")
    print("=" * 50)
    
    test_compat = TestBackwardCompatibility()
    await test_compat.test_anonymous_session_creation()
    await test_compat.test_anonymous_message_tracking()
    await test_compat.test_anonymous_search_tracking()
    await test_compat.test_existing_analytics_functions_work()
    print("✅ Backward compatibility tests passed")


async def run_performance_tests():
    """Run performance tests"""
    print("\n⚡ Running Performance Tests")
    print("=" * 50)
    
    test_perf = TestPerformance()
    await test_perf.test_user_id_query_performance()
    await test_perf.test_compound_index_performance()
    await test_perf.test_analytics_function_performance()
    print("✅ Performance tests passed")


async def main():
    """Run all comprehensive tests"""
    print("🚀 Starting Comprehensive User Authentication Tests")
    print("=" * 60)
    
    try:
        await run_unit_tests()
        await run_integration_tests()
        await run_analytics_tests()
        await run_compatibility_tests()
        await run_performance_tests()
        
        print("\n🎉 ALL TESTS PASSED!")
        print("User authentication feature is working correctly.")
        
    except Exception as e:
        print(f"\n❌ TEST FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    return True


if __name__ == "__main__":
    success = asyncio.run(main())
    exit(0 if success else 1)