test_postgres.py

#!/usr/bin/env python3
"""
Test script to verify PostgreSQL database functionality
"""

import psycopg2
import json
import logging
from datetime import datetime
import uuid

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Database configuration
POSTGRES_CONFIG = {
    'host': 'localhost',
    'port': 5432,
    'database': 'arrest_data',
    'user': 'arrest_user',
    'password': 'arrest_password'
}

def test_connection():
    """Test basic database connection"""
    try:
        conn = psycopg2.connect(**POSTGRES_CONFIG)
        cursor = conn.cursor()
        
        # Test basic query
        cursor.execute("SELECT version();")
        version = cursor.fetchone()
        logger.info(f"✅ Connected to PostgreSQL: {version[0]}")
        
        cursor.close()
        conn.close()
        return True
    except Exception as e:
        logger.error(f"❌ Connection test failed: {e}")
        return False

def test_schema():
    """Test if the arrest_data schema exists"""
    try:
        conn = psycopg2.connect(**POSTGRES_CONFIG)
        cursor = conn.cursor()
        
        # Check if schema exists
        cursor.execute("SELECT schema_name FROM information_schema.schemata WHERE schema_name = 'arrest_data';")
        schema = cursor.fetchone()
        
        if schema:
            logger.info("✅ arrest_data schema exists")
        else:
            logger.error("❌ arrest_data schema not found")
            return False
        
        cursor.close()
        conn.close()
        return True
    except Exception as e:
        logger.error(f"❌ Schema test failed: {e}")
        return False

def test_tables():
    """Test if all required tables exist"""
    try:
        conn = psycopg2.connect(**POSTGRES_CONFIG)
        cursor = conn.cursor()
        
        # Set search path
        cursor.execute("SET search_path TO arrest_data, public;")
        
        # Check if tables exist
        tables = ['clients', 'sessions', 'queries', 'messages']
        for table in tables:
            cursor.execute(f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = 'arrest_data' AND table_name = '{table}');")
            exists = cursor.fetchone()[0]
            if exists:
                logger.info(f"✅ Table '{table}' exists")
            else:
                logger.error(f"❌ Table '{table}' not found")
                return False
        
        cursor.close()
        conn.close()
        return True
    except Exception as e:
        logger.error(f"❌ Tables test failed: {e}")
        return False

def test_basic_operations():
    """Test basic CRUD operations"""
    try:
        conn = psycopg2.connect(**POSTGRES_CONFIG)
        cursor = conn.cursor()
        
        # Set search path
        cursor.execute("SET search_path TO arrest_data, public;")
        
        # Test client registration with unique data
        unique_suffix = str(uuid.uuid4())[:8]
        test_client = {
            'name': f'Test User {unique_suffix}',
            'nickname': f'testuser_{unique_suffix}',
            'email': f'test_{unique_suffix}@example.com',
            'password': 'testpassword'
        }
        
        # Insert test client
        cursor.execute("""
            INSERT INTO clients (name, nickname, email, password)
            VALUES (%s, %s, %s, %s)
            RETURNING id;
        """, (test_client['name'], test_client['nickname'], test_client['email'], test_client['password']))
        
        result = cursor.fetchone()
        if result:
            client_id = result[0]
            logger.info(f"✅ Test client created with ID: {client_id}")
        else:
            logger.error("❌ Failed to create test client")
            return False
        
        # Test session creation
        cursor.execute("""
            INSERT INTO sessions (client_id, address, start_time)
            VALUES (%s, %s, %s)
            RETURNING id;
        """, (client_id, '127.0.0.1', datetime.now().isoformat()))
        
        session_id = cursor.fetchone()[0]
        logger.info(f"✅ Test session created with ID: {session_id}")
        
        # Test query logging
        test_params = {'test': 'data', 'number': 42}
        cursor.execute("""
            INSERT INTO queries (client_id, session_id, query_type, parameters)
            VALUES (%s, %s, %s, %s)
            RETURNING id;
        """, (client_id, session_id, 'test_query', json.dumps(test_params)))
        
        query_id = cursor.fetchone()[0]
        logger.info(f"✅ Test query logged with ID: {query_id}")
        
        # Test message creation
        cursor.execute("""
            INSERT INTO messages (sender_type, sender_id, recipient_type, recipient_id, message)
            VALUES (%s, %s, %s, %s, %s)
            RETURNING id;
        """, ('client', client_id, 'all', 0, 'Test message'))
        
        message_id = cursor.fetchone()[0]
        logger.info(f"✅ Test message created with ID: {message_id}")
        
        # Test data retrieval
        cursor.execute("SELECT COUNT(*) FROM clients;")
        client_count = cursor.fetchone()[0]
        logger.info(f"✅ Total clients in database: {client_count}")
        
        cursor.execute("SELECT COUNT(*) FROM sessions;")
        session_count = cursor.fetchone()[0]
        logger.info(f"✅ Total sessions in database: {session_count}")
        
        cursor.execute("SELECT COUNT(*) FROM queries;")
        query_count = cursor.fetchone()[0]
        logger.info(f"✅ Total queries in database: {query_count}")
        
        cursor.execute("SELECT COUNT(*) FROM messages;")
        message_count = cursor.fetchone()[0]
        logger.info(f"✅ Total messages in database: {message_count}")
        
        # Clean up test data
        cursor.execute("DELETE FROM messages WHERE id = %s", (message_id,))
        cursor.execute("DELETE FROM queries WHERE id = %s", (query_id,))
        cursor.execute("DELETE FROM sessions WHERE id = %s", (session_id,))
        cursor.execute("DELETE FROM clients WHERE id = %s", (client_id,))
        
        conn.commit()
        cursor.close()
        conn.close()
        
        return True
    except Exception as e:
        logger.error(f"❌ Basic operations test failed: {e}")
        return False

def main():
    """Run all tests"""
    logger.info("🧪 Starting PostgreSQL database tests...")
    
    tests = [
        ("Connection", test_connection),
        ("Schema", test_schema),
        ("Tables", test_tables),
        ("Basic Operations", test_basic_operations)
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        logger.info(f"\n--- Testing {test_name} ---")
        if test_func():
            passed += 1
        else:
            logger.error(f"❌ {test_name} test failed")
    
    logger.info(f"\n📊 Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        logger.info("🎉 All tests passed! PostgreSQL database is working correctly.")
        return True
    else:
        logger.error("❌ Some tests failed. Please check the PostgreSQL setup.")
        return False

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