[model-gateway] rename py_test to e2e_test (#16454)

2026-01-04 22:24:34 -08:00
parent f02d82211a
commit e6fe092dcc
38 changed files with 9 additions and 9 deletions
@@ -0,0 +1,153 @@
+import datetime
+import os
+import socket
+import subprocess
+import sys
+import time
+from contextlib import closing
+
+import requests
+from cryptography import x509
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.x509.oid import NameOID
+from urllib3.exceptions import InsecureRequestWarning
+
+# Suppress insecure request warnings due to self-signed cert
+requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
+
+
+def find_free_port() -> int:
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+
+
+def generate_self_signed_cert(cert_path: str, key_path: str) -> None:
+    """Generate a self-signed certificate and private key for localhost."""
+    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+
+    subject = issuer = x509.Name(
+        [
+            x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
+            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "California"),
+            x509.NameAttribute(NameOID.LOCALITY_NAME, "San Francisco"),
+            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
+            x509.NameAttribute(NameOID.COMMON_NAME, "localhost"),
+        ]
+    )
+
+    cert = (
+        x509.CertificateBuilder()
+        .subject_name(subject)
+        .issuer_name(issuer)
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.datetime.utcnow())
+        .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=10))
+        .add_extension(
+            x509.SubjectAlternativeName([x509.DNSName("localhost")]), critical=False
+        )
+        .sign(key, hashes.SHA256())
+    )
+
+    with open(key_path, "wb") as f:
+        f.write(
+            key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.TraditionalOpenSSL,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+        )
+
+    with open(cert_path, "wb") as f:
+        f.write(cert.public_bytes(serialization.Encoding.PEM))
+
+
+def test_tls_server() -> None:
+    """End-to-end test for TLS-enabled router startup and basic endpoints."""
+    cert_path = "cert.pem"
+    key_path = "key.pem"
+    generate_self_signed_cert(cert_path, key_path)
+
+    port = find_free_port()
+
+    cmd = [
+        sys.executable,
+        "-m",
+        "sglang_router.launch_router",
+        "--worker-urls",
+        "http://127.0.0.1:9999",  # Dummy worker
+        "--host",
+        "127.0.0.1",
+        "--port",
+        str(port),
+        "--tls-cert-path",
+        cert_path,
+        "--tls-key-path",
+        key_path,
+        "--log-level",
+        "info",
+    ]
+
+    proc = subprocess.Popen(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+    )
+
+    try:
+        # Wait for server to start and respond to health check
+        start_time = time.time()
+        while time.time() - start_time < 15:
+            try:
+                response = requests.get(
+                    f"https://localhost:{port}/health", verify=False, timeout=2
+                )
+                if response.status_code == 200:
+                    break
+            except requests.RequestException:
+                pass
+
+            if proc.poll() is not None:
+                stdout, stderr = proc.communicate()
+                raise RuntimeError(
+                    f"Router process died early.\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}"
+                )
+
+            time.sleep(0.5)
+        else:
+            raise TimeoutError("Server did not become healthy within 15 seconds")
+
+        # Verify basic endpoints work over TLS
+        models_resp = requests.get(
+            f"https://localhost:{port}/v1/models", verify=False, timeout=2
+        )
+        assert models_resp.status_code in (
+            200,
+            503,
+        )  # 503 expected with no healthy workers
+
+        # Minimal generate request (should be rejected or queued)
+        gen_payload = {"model": "dummy", "prompt": "test", "max_new_tokens": 1}
+        gen_resp = requests.post(
+            f"https://localhost:{port}/generate",
+            json=gen_payload,
+            verify=False,
+            timeout=2,
+        )
+        assert gen_resp.status_code in (
+            200,
+            400,
+            503,
+        )  # Various valid responses with dummy worker
+
+    finally:
+        proc.terminate()
+        try:
+            proc.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+
+        for path in (cert_path, key_path):
+            if os.path.exists(path):
+                os.remove(path)