headscale-tailscale-replacement/test_headscale_models.py

#!/usr/bin/env python3
"""
Headscale Database Model Test Script
테스트를 위해 실제 SQLite DB에 연결하여 데이터 조회
"""

import sys
import os
from datetime import datetime, timedelta
from pathlib import Path

# Add current directory to path for importing models
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    from sqlalchemy import create_engine, text
    from sqlalchemy.orm import sessionmaker
    from headscale_models import (
        User, Node, PreAuthKey, ApiKey, Policy,
        PharmacyInfo, MachineSpecs, MonitoringData,
        create_all_tables
    )
    print("✅ SQLAlchemy models imported successfully")
except ImportError as e:
    print(f"❌ Failed to import models: {e}")
    print("💡 Install required packages: pip install sqlalchemy")
    sys.exit(1)


def test_database_connection():
    """데이터베이스 연결 테스트"""
    db_path = Path("data/db.sqlite")
    if not db_path.exists():
        print(f"❌ Database file not found: {db_path}")
        return None
    
    DATABASE_URL = f"sqlite:///{db_path}"
    print(f"🔗 Connecting to: {DATABASE_URL}")
    
    try:
        engine = create_engine(DATABASE_URL)
        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
        session = SessionLocal()
        
        # Test connection with a simple query
        result = session.execute(text("SELECT COUNT(*) FROM users")).scalar()
        print(f"✅ Database connection successful. User count: {result}")
        return session, engine
    except Exception as e:
        print(f"❌ Database connection failed: {e}")
        return None, None


def test_user_model(session):
    """User 모델 테스트"""
    print("\n" + "="*50)
    print("📊 TESTING USER MODEL")
    print("="*50)
    
    users = session.query(User).all()
    print(f"📋 Total users: {len(users)}")
    
    for user in users:
        print(f"\n👤 {user}")
        print(f"   - Created: {user.created_at}")
        print(f"   - Display Name: {user.display_name or 'Not set'}")
        print(f"   - Email: {user.email or 'Not set'}")
        print(f"   - Deleted: {user.is_deleted()}")
        print(f"   - Nodes Count: {len(user.nodes)}")


def test_node_model(session):
    """Node 모델 테스트"""
    print("\n" + "="*50)
    print("💻 TESTING NODE MODEL")
    print("="*50)
    
    nodes = session.query(Node).all()
    print(f"📋 Total nodes: {len(nodes)}")
    
    for node in nodes:
        print(f"\n🖥️  {node}")
        print(f"   - Given Name: {node.given_name}")
        print(f"   - User: {node.user.name if node.user else 'None'}")
        print(f"   - Online: {'🟢 Yes' if node.is_online() else '🔴 No'}")
        print(f"   - Last Seen: {node.last_seen}")
        print(f"   - Endpoints: {len(node.get_endpoints())} endpoint(s)")
        
        # Host info details
        host_info = node.get_host_info()
        if host_info:
            print(f"   - OS: {host_info.get('OS', 'Unknown')} {host_info.get('OSVersion', '')}")
            print(f"   - Hostname: {host_info.get('Hostname', 'Unknown')}")
            print(f"   - Machine: {host_info.get('Machine', 'Unknown')}")


def test_api_key_model(session):
    """API Key 모델 테스트"""
    print("\n" + "="*50)
    print("🔑 TESTING API KEY MODEL")
    print("="*50)
    
    api_keys = session.query(ApiKey).all()
    print(f"📋 Total API keys: {len(api_keys)}")
    
    for key in api_keys:
        print(f"\n🔐 {key}")
        print(f"   - Expired: {'❌ Yes' if key.is_expired() else '✅ No'}")
        print(f"   - Created: {key.created_at}")
        print(f"   - Expires: {key.expiration}")
        print(f"   - Last Used: {key.last_seen or 'Never'}")


def test_pre_auth_key_model(session):
    """Pre-Auth Key 모델 테스트"""
    print("\n" + "="*50)
    print("🎫 TESTING PRE-AUTH KEY MODEL")
    print("="*50)
    
    pre_auth_keys = session.query(PreAuthKey).all()
    print(f"📋 Total pre-auth keys: {len(pre_auth_keys)}")
    
    for key in pre_auth_keys:
        print(f"\n🎟️  {key}")
        print(f"   - User: {key.user.name if key.user else 'None'}")
        print(f"   - Reusable: {'✅ Yes' if key.reusable else '❌ No'}")
        print(f"   - Used: {'✅ Yes' if key.used else '❌ No'}")
        print(f"   - Valid: {'✅ Yes' if key.is_valid() else '❌ No'}")
        print(f"   - Expires: {key.expiration}")
        print(f"   - Tags: {key.get_tags()}")


def test_policy_model(session):
    """Policy 모델 테스트"""
    print("\n" + "="*50)
    print("📜 TESTING POLICY MODEL")
    print("="*50)
    
    policies = session.query(Policy).all()
    print(f"📋 Total policies: {len(policies)}")
    
    for policy in policies:
        print(f"\n📄 {policy}")
        policy_data = policy.get_policy_data()
        if policy_data:
            print(f"   - ACL Rules: {len(policy_data.get('acls', []))}")
            print(f"   - Groups: {len(policy_data.get('groups', {}))}")


def create_sample_extended_data(session, engine):
    """확장 테이블용 샘플 데이터 생성"""
    print("\n" + "="*50)
    print("🏥 CREATING SAMPLE PHARMACY DATA")
    print("="*50)
    
    # Create extended tables
    create_all_tables(engine)
    
    # Get first user
    user = session.query(User).first()
    if not user:
        print("❌ No users found. Cannot create pharmacy info.")
        return
    
    # Check if pharmacy info already exists
    existing_pharmacy = session.query(PharmacyInfo).filter_by(user_id=user.name).first()
    if existing_pharmacy:
        print(f"ℹ️  Pharmacy info already exists for user '{user.name}'")
        return
    
    # Create pharmacy info
    pharmacy = PharmacyInfo(
        user_id=user.name,
        pharmacy_name="서울중앙약국",
        business_number="123-45-67890",
        address="서울시 강남구 테헤란로 123",
        phone="02-1234-5678",
        manager_name="홍길동",
        proxmox_host="192.168.1.100",
        proxmox_api_token="sample_token_here"
    )
    session.add(pharmacy)
    
    # Get first node
    node = session.query(Node).first()
    if node:
        # Create machine specs
        specs = MachineSpecs(
            machine_id=node.id,
            pharmacy_id=1,  # Will be set properly after pharmacy is committed
            cpu_model="Intel Core i7-12700",
            cpu_cores=12,
            ram_gb=32,
            storage_gb=1000,
            gpu_model="NVIDIA GTX 1660"
        )
        session.add(specs)
        
        # Create monitoring data
        monitoring = MonitoringData(
            machine_id=node.id,
            cpu_usage="75.50",
            memory_usage="60.25",
            disk_usage="45.00",
            cpu_temperature=65,
            network_rx_bytes=1024000,
            network_tx_bytes=512000,
            vm_count=5,
            vm_running=4
        )
        session.add(monitoring)
    
    try:
        session.commit()
        print("✅ Sample extended data created successfully")
    except Exception as e:
        session.rollback()
        print(f"❌ Failed to create sample data: {e}")


def test_extended_models(session):
    """확장된 모델 테스트"""
    print("\n" + "="*50)
    print("🏥 TESTING EXTENDED MODELS (FARMQ)")
    print("="*50)
    
    # Test pharmacy info
    pharmacies = session.query(PharmacyInfo).all()
    print(f"🏪 Total pharmacies: {len(pharmacies)}")
    for pharmacy in pharmacies:
        print(f"   - {pharmacy}")
    
    # Test machine specs
    specs = session.query(MachineSpecs).all()
    print(f"⚙️  Total machine specs: {len(specs)}")
    for spec in specs:
        print(f"   - {spec}")
    
    # Test monitoring data
    monitoring = session.query(MonitoringData).all()
    print(f"📊 Total monitoring records: {len(monitoring)}")
    for monitor in monitoring:
        print(f"   - {monitor}")


def main():
    """메인 테스트 함수"""
    print("🧪 HEADSCALE DATABASE MODEL TEST")
    print("=" * 60)
    
    # Connect to database
    session, engine = test_database_connection()
    if not session:
        return
    
    try:
        # Test core models
        test_user_model(session)
        test_node_model(session)
        test_api_key_model(session)
        test_pre_auth_key_model(session)
        test_policy_model(session)
        
        # Create sample extended data (if needed)
        create_sample_extended_data(session, engine)
        
        # Test extended models
        test_extended_models(session)
        
        print("\n" + "="*60)
        print("🎉 ALL TESTS COMPLETED SUCCESSFULLY!")
        print("="*60)
        
    except Exception as e:
        print(f"\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()
        
    finally:
        session.close()


if __name__ == "__main__":
    main()