test_postgresql.py

import json
import subprocess
import sys
import time

# Database credentials (replace with your test DB)
PG_DB_CONFIG = {
    "db_type": "postgresql",
    "host": "localhost",
    "port": 5432,
    "user": "testuser",
    "password": "testpassword",
    "database": "testdb",
    "connection_name": "pg_test_conn"
}

_process = None

def start_server():
    global _process
    _process = subprocess.Popen(
        [sys.executable, "mcp_server.py"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=0
    )
    # Wait for server to be ready
    time.sleep(1)
    return _process

def send_request(request):
    request_json = json.dumps(request) + "\n"
    _process.stdin.write(request_json)
    _process.stdin.flush()

    response_line = _process.stdout.readline()
    if response_line.strip():
        return json.loads(response_line.strip())
    return None

def send_notification(notification):
    notification_json = json.dumps(notification) + "\n"
    _process.stdin.write(notification_json)
    _process.stdin.flush()

def cleanup():
    if _process:
        _process.terminate()
        _process.wait(timeout=5)

def test_postgresql_connection():
    print("🚀 Testing PostgreSQL Connection...")

    # 1. Initialize
    send_request({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}})
    send_notification({"jsonrpc": "2.0", "method": "notifications/initialized"})

    # 2. Connect to PostgreSQL
    connect_req = {
        "jsonrpc": "2.0",
        "id": 2,
        "method": "tools/call",
        "params": {
            "name": "connect_database",
            "arguments": PG_DB_CONFIG
        }
    }

    response = send_request(connect_req)
    content = json.loads(response['result']['content'][0]['text'])

    if not content.get("success"):
        print(f"❌ Connection failed: {content.get('error')}")
        return False
    print("✅ Connection successful!")

    # 3. List tables
    list_tables_req = {
        "jsonrpc": "2.0",
        "id": 3,
        "method": "tools/call",
        "params": {
            "name": "list_tables",
            "arguments": {"connection_name": PG_DB_CONFIG["connection_name"]}
        }
    }
    response = send_request(list_tables_req)
    content = json.loads(response['result']['content'][0]['text'])
    print(f"✅ List tables result: {content}")

    return True

if __name__ == "__main__":
    start_server()
    try:
        if not test_postgresql_connection():
            sys.exit(1)
    finally:
        cleanup()
        print("✅ PostgreSQL tests finished.")