with_server.py 3.6 KB

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