1#!/usr/bin/env python3
  2"""
  3Start one or more servers, wait for them to be ready, run a command, then clean up.
  4
  5Usage:
  6    # Single server
  7    python scripts/with_server.py --server "npm run dev" --port 5173 -- python automation.py
  8    python scripts/with_server.py --server "npm start" --port 3000 -- python test.py
  9
 10    # Multiple servers
 11    python scripts/with_server.py \
 12      --server "cd backend && python server.py" --port 3000 \
 13      --server "cd frontend && npm run dev" --port 5173 \
 14      -- python test.py
 15"""
 16
 17import subprocess
 18import socket
 19import time
 20import sys
 21import argparse
 22
 23def is_server_ready(port, timeout=30):
 24    """Wait for server to be ready by polling the port."""
 25    start_time = time.time()
 26    while time.time() - start_time < timeout:
 27        try:
 28            with socket.create_connection(('localhost', port), timeout=1):
 29                return True
 30        except (socket.error, ConnectionRefusedError):
 31            time.sleep(0.5)
 32    return False
 33
 34
 35def main():
 36    parser = argparse.ArgumentParser(description='Run command with one or more servers')
 37    parser.add_argument('--server', action='append', dest='servers', required=True, help='Server command (can be repeated)')
 38    parser.add_argument('--port', action='append', dest='ports', type=int, required=True, help='Port for each server (must match --server count)')
 39    parser.add_argument('--timeout', type=int, default=30, help='Timeout in seconds per server (default: 30)')
 40    parser.add_argument('command', nargs=argparse.REMAINDER, help='Command to run after server(s) ready')
 41
 42    args = parser.parse_args()
 43
 44    # Remove the '--' separator if present
 45    if args.command and args.command[0] == '--':
 46        args.command = args.command[1:]
 47
 48    if not args.command:
 49        print("Error: No command specified to run")
 50        sys.exit(1)
 51
 52    # Parse server configurations
 53    if len(args.servers) != len(args.ports):
 54        print("Error: Number of --server and --port arguments must match")
 55        sys.exit(1)
 56
 57    servers = []
 58    for cmd, port in zip(args.servers, args.ports):
 59        servers.append({'cmd': cmd, 'port': port})
 60
 61    server_processes = []
 62
 63    try:
 64        # Start all servers
 65        for i, server in enumerate(servers):
 66            print(f"Starting server {i+1}/{len(servers)}: {server['cmd']}")
 67
 68            # Use shell=True to support commands with cd and &&
 69            process = subprocess.Popen(
 70                server['cmd'],
 71                shell=True,
 72                stdout=subprocess.PIPE,
 73                stderr=subprocess.PIPE
 74            )
 75            server_processes.append(process)
 76
 77            # Wait for this server to be ready
 78            print(f"Waiting for server on port {server['port']}...")
 79            if not is_server_ready(server['port'], timeout=args.timeout):
 80                raise RuntimeError(f"Server failed to start on port {server['port']} within {args.timeout}s")
 81
 82            print(f"Server ready on port {server['port']}")
 83
 84        print(f"\nAll {len(servers)} server(s) ready")
 85
 86        # Run the command
 87        print(f"Running: {' '.join(args.command)}\n")
 88        result = subprocess.run(args.command)
 89        sys.exit(result.returncode)
 90
 91    finally:
 92        # Clean up all servers
 93        print(f"\nStopping {len(server_processes)} server(s)...")
 94        for i, process in enumerate(server_processes):
 95            try:
 96                process.terminate()
 97                process.wait(timeout=5)
 98            except subprocess.TimeoutExpired:
 99                process.kill()
100                process.wait()
101            print(f"Server {i+1} stopped")
102        print("All servers stopped")
103
104
105if __name__ == '__main__':
106    main()