78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Small Switchyard-compatible health/models stand-in for integration tests."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import signal
|
||
|
|
import tomllib
|
||
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
def route_ids(path: Path) -> list[str]:
|
||
|
|
with path.open("rb") as handle:
|
||
|
|
data = tomllib.load(handle)
|
||
|
|
return [str(value["id"]) for value in data.get("routes", {}).values()]
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--version", action="version", version="switchyard-server 0.2.0")
|
||
|
|
parser.add_argument("--config", type=Path, required=True)
|
||
|
|
parser.add_argument("--host", default="127.0.0.1")
|
||
|
|
parser.add_argument("--port", type=int, default=4000)
|
||
|
|
parser.add_argument("--routing-log-file")
|
||
|
|
parser.add_argument("--dry-run", action="store_true")
|
||
|
|
args = parser.parse_args()
|
||
|
|
models = route_ids(args.config)
|
||
|
|
if args.dry_run:
|
||
|
|
print(json.dumps({"valid": True, "routes": models}))
|
||
|
|
return 0
|
||
|
|
|
||
|
|
class Handler(BaseHTTPRequestHandler):
|
||
|
|
def log_message(self, _format: str, *_args: object) -> None:
|
||
|
|
return
|
||
|
|
|
||
|
|
def respond(self, status: int, value: object) -> None:
|
||
|
|
body = json.dumps(value).encode()
|
||
|
|
self.send_response(status)
|
||
|
|
self.send_header("Content-Type", "application/json")
|
||
|
|
self.send_header("Content-Length", str(len(body)))
|
||
|
|
self.end_headers()
|
||
|
|
self.wfile.write(body)
|
||
|
|
|
||
|
|
def do_GET(self) -> None: # noqa: N802
|
||
|
|
if self.path == "/health":
|
||
|
|
self.respond(200, {"status": "ok"})
|
||
|
|
elif self.path in {"/v1/models", "/models"}:
|
||
|
|
self.respond(
|
||
|
|
200,
|
||
|
|
{
|
||
|
|
"object": "list",
|
||
|
|
"data": [{"id": item, "object": "model"} for item in models],
|
||
|
|
},
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
self.respond(404, {"error": "not found"})
|
||
|
|
|
||
|
|
server = ThreadingHTTPServer((args.host, args.port), Handler)
|
||
|
|
|
||
|
|
def stop(_signum: int, _frame: object) -> None:
|
||
|
|
raise KeyboardInterrupt
|
||
|
|
|
||
|
|
signal.signal(signal.SIGTERM, stop)
|
||
|
|
signal.signal(signal.SIGINT, stop)
|
||
|
|
try:
|
||
|
|
server.serve_forever(poll_interval=0.1)
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
pass
|
||
|
|
finally:
|
||
|
|
server.server_close()
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|