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()