atlas/home-voice-agent/mcp-server/server/test_admin_api.py

#!/usr/bin/env python3
"""
Tests for Admin API endpoints.
"""

import sys
from pathlib import Path
import tempfile
import sqlite3
import json
from datetime import datetime

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

try:
    from fastapi.testclient import TestClient
    from fastapi import FastAPI
    from server.admin_api import router
except ImportError as e:
    print(f"⚠️  Import error: {e}")
    print("   Install dependencies: cd mcp-server && pip install -r requirements.txt")
    sys.exit(1)

# Create test app
app = FastAPI()
app.include_router(router)

client = TestClient(app)

# Test data directory
TEST_DATA_DIR = Path(__file__).parent.parent.parent / "data" / "test_admin"
TEST_DATA_DIR.mkdir(parents=True, exist_ok=True)


def setup_test_databases():
    """Create test databases."""
    tokens_db = TEST_DATA_DIR / "tokens.db"
    tokens_db.parent.mkdir(parents=True, exist_ok=True)
    
    if tokens_db.exists():
        tokens_db.unlink()
    
    conn = sqlite3.connect(str(tokens_db))
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE revoked_tokens (
            token_id TEXT PRIMARY KEY,
            device_id TEXT,
            revoked_at TEXT NOT NULL,
            reason TEXT,
            revoked_by TEXT
        )
    """)
    cursor.execute("""
        CREATE TABLE devices (
            device_id TEXT PRIMARY KEY,
            name TEXT,
            last_seen TEXT,
            status TEXT DEFAULT 'active',
            created_at TEXT NOT NULL
        )
    """)
    cursor.execute("""
        INSERT INTO devices (device_id, name, last_seen, status, created_at)
        VALUES ('device-1', 'Test Device', '2026-01-01T00:00:00', 'active', '2026-01-01T00:00:00')
    """)
    conn.commit()
    conn.close()
    
    # Logs directory
    logs_dir = TEST_DATA_DIR / "logs"
    logs_dir.mkdir(exist_ok=True)
    
    # Create test log file
    log_file = logs_dir / "llm_2026-01-01.log"
    log_file.write_text(json.dumps({
        "timestamp": "2026-01-01T00:00:00",
        "level": "INFO",
        "agent_type": "family",
        "tool_calls": ["get_current_time"],
        "message": "Test log entry"
    }) + "\n")
    
    return {
        "tokens": tokens_db,
        "logs": logs_dir
    }


def test_enhanced_logs():
    """Test /api/admin/logs/enhanced endpoint."""
    import server.admin_api as admin_api
    original_logs = admin_api.LOGS_DIR
    
    try:
        test_dbs = setup_test_databases()
        admin_api.LOGS_DIR = test_dbs["logs"]
        
        response = client.get("/api/admin/logs/enhanced?limit=10")
        assert response.status_code == 200
        data = response.json()
        assert "logs" in data
        assert "total" in data
        assert len(data["logs"]) >= 1
        
        # Test filters
        response = client.get("/api/admin/logs/enhanced?level=INFO&agent_type=family")
        assert response.status_code == 200
        
        print("✅ Enhanced logs endpoint test passed")
        return True
    finally:
        admin_api.LOGS_DIR = original_logs


def test_revoke_token():
    """Test /api/admin/revoke_token endpoint."""
    import server.admin_api as admin_api
    original_tokens = admin_api.TOKENS_DB
    
    try:
        test_dbs = setup_test_databases()
        admin_api.TOKENS_DB = test_dbs["tokens"]
        admin_api._init_tokens_db()
        
        response = client.post(
            "/api/admin/revoke_token",
            json={
                "token_id": "test-token-1",
                "reason": "Test revocation",
                "revoked_by": "admin"
            }
        )
        assert response.status_code == 200
        data = response.json()
        assert data["success"] is True
        
        # Verify token is in database
        conn = sqlite3.connect(str(test_dbs["tokens"]))
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM revoked_tokens WHERE token_id = ?", ("test-token-1",))
        row = cursor.fetchone()
        assert row is not None
        conn.close()
        
        print("✅ Revoke token endpoint test passed")
        return True
    finally:
        admin_api.TOKENS_DB = original_tokens


def test_list_revoked_tokens():
    """Test /api/admin/list_revoked_tokens endpoint."""
    import server.admin_api as admin_api
    original_tokens = admin_api.TOKENS_DB
    
    try:
        test_dbs = setup_test_databases()
        admin_api.TOKENS_DB = test_dbs["tokens"]
        admin_api._init_tokens_db()
        
        # Add a revoked token first
        conn = sqlite3.connect(str(test_dbs["tokens"]))
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO revoked_tokens (token_id, device_id, revoked_at, reason, revoked_by)
            VALUES ('test-token-2', 'device-1', '2026-01-01T00:00:00', 'Test', 'admin')
        """)
        conn.commit()
        conn.close()
        
        response = client.get("/api/admin/list_revoked_tokens")
        assert response.status_code == 200
        data = response.json()
        assert "tokens" in data
        assert len(data["tokens"]) >= 1
        
        print("✅ List revoked tokens endpoint test passed")
        return True
    finally:
        admin_api.TOKENS_DB = original_tokens


def test_register_device():
    """Test /api/admin/register_device endpoint."""
    import server.admin_api as admin_api
    original_tokens = admin_api.TOKENS_DB
    
    try:
        test_dbs = setup_test_databases()
        admin_api.TOKENS_DB = test_dbs["tokens"]
        admin_api._init_tokens_db()
        
        response = client.post(
            "/api/admin/register_device",
            json={
                "device_id": "test-device-2",
                "name": "Test Device 2"
            }
        )
        assert response.status_code == 200
        data = response.json()
        assert data["success"] is True
        
        # Verify device is in database
        conn = sqlite3.connect(str(test_dbs["tokens"]))
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM devices WHERE device_id = ?", ("test-device-2",))
        row = cursor.fetchone()
        assert row is not None
        conn.close()
        
        print("✅ Register device endpoint test passed")
        return True
    finally:
        admin_api.TOKENS_DB = original_tokens


def test_list_devices():
    """Test /api/admin/list_devices endpoint."""
    import server.admin_api as admin_api
    original_tokens = admin_api.TOKENS_DB
    
    try:
        test_dbs = setup_test_databases()
        admin_api.TOKENS_DB = test_dbs["tokens"]
        admin_api._init_tokens_db()
        
        response = client.get("/api/admin/list_devices")
        assert response.status_code == 200
        data = response.json()
        assert "devices" in data
        assert len(data["devices"]) >= 1
        
        print("✅ List devices endpoint test passed")
        return True
    finally:
        admin_api.TOKENS_DB = original_tokens


def test_revoke_device():
    """Test /api/admin/revoke_device endpoint."""
    import server.admin_api as admin_api
    original_tokens = admin_api.TOKENS_DB
    
    try:
        test_dbs = setup_test_databases()
        admin_api.TOKENS_DB = test_dbs["tokens"]
        admin_api._init_tokens_db()
        
        response = client.post(
            "/api/admin/revoke_device",
            json={
                "device_id": "device-1",
                "reason": "Test revocation"
            }
        )
        assert response.status_code == 200
        data = response.json()
        assert data["success"] is True
        
        # Verify device status is revoked
        conn = sqlite3.connect(str(test_dbs["tokens"]))
        cursor = conn.cursor()
        cursor.execute("SELECT status FROM devices WHERE device_id = ?", ("device-1",))
        row = cursor.fetchone()
        assert row is not None
        assert row[0] == "revoked"
        conn.close()
        
        print("✅ Revoke device endpoint test passed")
        return True
    finally:
        admin_api.TOKENS_DB = original_tokens


if __name__ == "__main__":
    print("=" * 60)
    print("Admin API Test Suite")
    print("=" * 60)
    print()
    
    try:
        test_enhanced_logs()
        test_revoke_token()
        test_list_revoked_tokens()
        test_register_device()
        test_list_devices()
        test_revoke_device()
        
        print()
        print("=" * 60)
        print("✅ All Admin API tests passed!")
        print("=" * 60)
    except Exception as e:
        print(f"\n❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)