prain: skeleton

tests/__init__.py

@@ -0,0 +1 @@
+

tests/helpers.py

@@ -0,0 +1,98 @@
+import json
+import os
+import tempfile
+import threading
+import time
+import unittest
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+
+class TmpCase(unittest.TestCase):
+    """A test case with a fresh temporary directory in self.dir."""
+
+    def setUp(self):
+        tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(tmp.cleanup)
+        self.dir = tmp.name
+
+    def write(self, name, text):
+        path = os.path.join(self.dir, name)
+        with open(path, "w") as f:
+            f.write(text)
+        return path
+
+
+def wait_for(predicate, timeout=5.0):
+    end = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() > end:
+            raise AssertionError("timed out waiting")
+        time.sleep(0.01)
+
+
+class StubVllm:
+    """A fake OpenAI-compatible server on a free port. Records every POST body. Streamed
+    chats answer "Hel", "lo", "!"; others answer "ok". Set .delay (seconds, or a function of
+    the request body) to make answers slow, .pause (the same) to stream the pieces that far
+    apart, and .fail to an HTTP code to reject chats."""
+
+    def __init__(self):
+        self.bodies = []
+        self.delay = 0.0
+        self.pause = 0.0
+        self.fail = None
+        stub = self
+
+        class Handler(BaseHTTPRequestHandler):
+            def do_GET(self):
+                self.send(200, "application/json", {"data": [{"id": "stub-model"}]})
+
+            def do_POST(self):
+                body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+                stub.bodies.append(body)
+                time.sleep(stub.delay(body) if callable(stub.delay) else stub.delay)
+                if stub.fail:
+                    return self.send(stub.fail, "application/json", {"error": "rejected"})
+                if not body.get("stream"):
+                    return self.send(200, "application/json",
+                                     {"choices": [{"message": {"content": "ok"}}]})
+                events = [{"choices": [{"delta": {"content": p}}]} for p in ("Hel", "lo", "!")]
+                pause = stub.pause(body) if callable(stub.pause) else stub.pause
+                if pause:  # stream piece by piece; HTTP/1.0 ends the body by closing
+                    return self.trickle([f"data: {json.dumps(e)}\n\n" for e in events], pause)
+                text = "".join(f"data: {json.dumps(e)}\n\n" for e in events) + "data: [DONE]\n\n"
+                self.send(200, "text/event-stream", text)
+
+            def trickle(self, chunks, pause):
+                try:
+                    self.send_response(200)
+                    self.send_header("Content-Type", "text/event-stream")
+                    self.end_headers()
+                    for i, chunk in enumerate(chunks + ["data: [DONE]\n\n"]):
+                        time.sleep(pause if i else 0)
+                        self.wfile.write(chunk.encode())
+                        self.wfile.flush()
+                except (BrokenPipeError, ConnectionResetError):
+                    pass  # the client aborted
+
+            def send(self, code, kind, payload):
+                data = (payload if isinstance(payload, str) else json.dumps(payload)).encode()
+                try:
+                    self.send_response(code)
+                    self.send_header("Content-Type", kind)
+                    self.send_header("Content-Length", str(len(data)))
+                    self.end_headers()
+                    self.wfile.write(data)
+                except (BrokenPipeError, ConnectionResetError):
+                    pass  # the client aborted
+
+            def log_message(self, *args):
+                pass
+
+        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+        self.url = f"http://127.0.0.1:{self.server.server_port}"
+        threading.Thread(target=self.server.serve_forever, daemon=True).start()
+
+    def close(self):
+        self.server.shutdown()
+        self.server.server_close()