test_persistence_fork.py

#!/usr/bin/env python3
"""
Test Persistence and Fork Features

This script tests:
1. Auto-save functionality in tutor_engine
2. Fork API endpoint
3. Session persistence to database
"""

import os
import sys
import json
import asyncio
from datetime import datetime

# Add backend to path
sys.path.insert(0, '/root/AIFeyman/backend')

from sqlmodel import SQLModel, create_engine, Session, text, select
from app.models.schema import User, LearningSessionDB, ChatMessageDB, LearningSessionCreate, ForkSessionRequest
from app.api.auth import get_password_hash
from app.core.auth import AuthConfig
from app.db.database import get_session

# Set test environment
os.environ["NEXTAUTH_SECRET"] = "test-secret-key"

def test_database_connection():
    """Test database connection and basic operations"""
    print("\n" + "="*60)
    print("🗄️  Testing Database Connection")
    print("="*60)

    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+psycopg2://postgres:postgres@127.0.0.1:5432/aifeyman"
    )

    engine = create_engine(database_url, echo=False)

    try:
        with Session(engine) as session:
            # Test connection
            result = session.exec(text("SELECT 1")).first()
            print(f"✅ Database connection successful")

            # Count users
            user_count = session.exec(text("SELECT COUNT(*) FROM user")).first()
            print(f"✅ Users in database: {user_count[0]}")

            # Count sessions
            session_count = session.exec(text("SELECT COUNT(*) FROM learningsessiondb")).first()
            print(f"✅ Sessions in database: {session_count[0]}")

            # Count messages
            message_count = session.exec(text("SELECT COUNT(*) FROM chatmessagedb")).first()
            print(f"✅ Messages in database: {message_count[0]}")

            return True

    except Exception as e:
        print(f"❌ Database connection failed: {e}")
        return False

def test_create_user():
    """Test creating a new user"""
    print("\n" + "="*60)
    print("👤 Testing User Creation")
    print("="*60)

    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+psycopg2://postgres:postgres@127.0.0.1:5432/aifeyman"
    )

    engine = create_engine(database_url, echo=False)

    try:
        with Session(engine) as session:
            # Check if test user exists
            existing_user_id = session.exec(
                select(User.id).where(User.email == "test@example.com")
            ).first()

            if existing_user_id:
                existing_user = session.get(User, existing_user_id)
                print(f"✅ Test user already exists: {existing_user.email}")
                return existing_user

            # Create new user
            hashed_password = get_password_hash("test123456")

            new_user = User(
                email="test@example.com",
                name="Test User",
                hashed_password=hashed_password
            )

            session.add(new_user)
            session.commit()
            session.refresh(new_user)

            print(f"✅ Created test user: {new_user.email}")
            return new_user

    except Exception as e:
        print(f"❌ User creation failed: {e}")
        return None

def test_create_session(user: User):
    """Test creating a learning session"""
    print("\n" + "="*60)
    print("📚 Testing Session Creation")
    print("="*60)

    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+psycopg2://postgres:postgres@127.0.0.1:5432/aifeyman"
    )

    engine = create_engine(database_url, echo=False)

    try:
        with Session(engine) as session:
            # Check if test session exists
            existing_session_id = session.exec(
                select(LearningSessionDB.id).where(
                    LearningSessionDB.user_id == user.id,
                    LearningSessionDB.topic == "测试主题"
                )
            ).first()

            if existing_session_id:
                existing_session = session.get(LearningSessionDB, existing_session_id)
                print(f"✅ Test session already exists: {existing_session.id}")
                return existing_session

            # Create new session
            test_session = LearningSessionDB(
                user_id=user.id,
                topic="测试主题",
                source_text="这是一个测试学习主题，用于验证持久化和 Fork 功能。",
                graph_data={
                    "nodes": [
                        {
                            "id": "concept-1",
                            "label": "概念1",
                            "status": "neutral",
                            "position": {"x": 100, "y": 100},
                            "data": {"description": "测试概念描述", "difficulty": 3}
                        }
                    ],
                    "edges": []
                },
                is_public=True
            )

            session.add(test_session)
            session.commit()
            session.refresh(test_session)

            print(f"✅ Created test session: {test_session.id}")
            print(f"   Topic: {test_session.topic}")
            print(f"   Is Public: {test_session.is_public}")
            print(f"   Graph Nodes: {len(test_session.graph_data.get('nodes', []))}")

            return test_session

    except Exception as e:
        print(f"❌ Session creation failed: {e}")
        import traceback
        traceback.print_exc()
        return None

def test_add_messages(session_db: LearningSessionDB):
    """Test adding chat messages to a session"""
    print("\n" + "="*60)
    print("💬 Testing Chat Messages")
    print("="*60)

    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+psycopg2://postgres:postgres@127.0.0.1:5432/aifeyman"
    )

    engine = create_engine(database_url, echo=False)

    try:
        with Session(engine) as session:
            # Check if messages exist
            message_count = session.exec(
                select(ChatMessageDB).where(ChatMessageDB.session_id == session_db.id)
            ).all()

            if message_count:
                print(f"✅ Session already has {len(message_count)} messages")
                return True

            # Add user message
            user_message = ChatMessageDB(
                session_id=session_db.id,
                role="user",
                content="我想学习这个主题"
            )
            session.add(user_message)

            # Add AI response
            ai_message = ChatMessageDB(
                session_id=session_db.id,
                role="assistant",
                content="很好！让我们开始学习第一个概念。"
            )
            session.add(ai_message)

            session.commit()

            print(f"✅ Added 2 messages to session")
            return True

    except Exception as e:
        print(f"❌ Adding messages failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_fork_session(source_session: LearningSessionDB, user: User):
    """Test forking a session"""
    print("\n" + "="*60)
    print("🔀 Testing Session Fork")
    print("="*60)

    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+psycopg2://postgres:postgres@127.0.0.1:5432/aifeyman"
    )

    engine = create_engine(database_url, echo=False)

    try:
        with Session(engine) as session:
            # Check if fork already exists
            fork_exists = session.exec(
                select(LearningSessionDB).where(
                    LearningSessionDB.fork_from_id == source_session.id,
                    LearningSessionDB.user_id == user.id
                )
            ).first()

            if fork_exists:
                print(f"✅ Fork already exists: {fork_exists.id}")
                return fork_exists

            # Create fork
            forked_session = LearningSessionDB(
                user_id=user.id,
                topic=f"{source_session.topic} (副本)",
                source_text=source_session.source_text,
                graph_data=source_session.graph_data.copy(),  # Snapshot clone
                is_public=False,
                fork_from_id=source_session.id
            )

            session.add(forked_session)
            session.commit()
            session.refresh(forked_session)

            print(f"✅ Created fork: {forked_session.id}")
            print(f"   Original Session: {source_session.id}")
            print(f"   Forked Session: {forked_session.id}")
            print(f"   Fork From ID: {forked_session.fork_from_id}")
            print(f"   Topic: {forked_session.topic}")
            print(f"   Is Public: {forked_session.is_public}")
            print(f"   Graph Nodes: {len(forked_session.graph_data.get('nodes', []))}")

            # Verify it's a snapshot clone (no chat history)
            fork_messages = session.exec(
                select(ChatMessageDB).where(ChatMessageDB.session_id == forked_session.id)
            ).all()

            print(f"   Fork Messages: {len(fork_messages)} (should be 0 - snapshot clone)")

            if len(fork_messages) == 0:
                print("✅ Verified: Fork is a snapshot clone (no chat history)")
            else:
                print("⚠️  Warning: Fork has chat history (should be empty)")

            return forked_session

    except Exception as e:
        print(f"❌ Fork failed: {e}")
        import traceback
        traceback.print_exc()
        return None

def test_session_list(user: User):
    """Test listing user sessions"""
    print("\n" + "="*60)
    print("📋 Testing Session List")
    print("="*60)

    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+psycopg2://postgres:postgres@127.0.0.1:5432/aifeyman"
    )

    engine = create_engine(database_url, echo=False)

    try:
        with Session(engine) as session:
            # Get all user sessions
            user_sessions = session.exec(
                select(LearningSessionDB).where(LearningSessionDB.user_id == user.id)
            ).all()

            print(f"✅ User has {len(user_sessions)} sessions:")
            for s in user_sessions:
                fork_status = f" (fork of {s.fork_from_id})" if s.fork_from_id else ""
                print(f"   - {s.id}: {s.topic}{fork_status}")

            return user_sessions

    except Exception as e:
        print(f"❌ Session list failed: {e}")
        return []

def main():
    """Run all tests"""
    print("\n" + "="*60)
    print("🔬 Persistence and Fork Feature Tests")
    print("="*60)

    # Test 1: Database Connection
    if not test_database_connection():
        print("\n❌ Cannot proceed without database connection")
        print("Please start the database: /root/AIFeyman/start_db.sh")
        sys.exit(1)

    # Test 2: Create User
    user = test_create_user()
    if not user:
        print("\n❌ Cannot proceed without user")
        sys.exit(1)

    # Test 3: Create Session
    session_db = test_create_session(user)
    if not session_db:
        print("\n❌ Cannot proceed without session")
        sys.exit(1)

    # Test 4: Add Messages
    if not test_add_messages(session_db):
        print("\n⚠️  Warning: Failed to add messages")

    # Test 5: Fork Session
    fork = test_fork_session(session_db, user)
    if not fork:
        print("\n⚠️  Warning: Fork test failed")

    # Test 6: List Sessions
    sessions = test_session_list(user)

    # Summary
    print("\n" + "="*60)
    print("📊 Test Summary")
    print("="*60)
    print(f"✅ Database Connection: OK")
    print(f"✅ User Creation: OK")
    print(f"✅ Session Creation: OK")
    print(f"{'✅' if test_add_messages(session_db) else '⚠️'} Chat Messages: {'OK' if test_add_messages(session_db) else 'FAILED'}")
    print(f"{'✅' if fork else '⚠️'} Session Fork: {'OK' if fork else 'FAILED'}")
    print(f"✅ Session List: {len(sessions)} sessions")
    print("="*60)

    if fork:
        print("\n✅ All core features working!")
        print("\n📝 Next steps:")
        print("1. Test auto-save in tutor_engine (requires LLM integration)")
        print("2. Test fork API via HTTP endpoint")
        print("3. Test loading session from database")
    else:
        print("\n⚠️  Some tests failed - check errors above")

    print("="*60)

if __name__ == "__main__":
    main()