poll for events from test server

tests/llmobs/conftest.py

@@ -1,7 +1,8 @@
 from http.server import BaseHTTPRequestHandler
 from http.server import HTTPServer
+import json
 import os
-import socket
+import time
 import threading
 
 import mock
@@ -197,48 +198,61 @@ def llmobs_span_writer(_llmobs_backend):
     yield sw
 
 
-class LLMObsBackend(BaseHTTPRequestHandler):
+class LLMObsServer(BaseHTTPRequestHandler):
+    """A mock server for the LLMObs backend used to capture the requests made by the client.
+
+    Python's HTTPRequestHandler is a bit weird and uses a class rather than an instance
+    for running an HTTP server so the requests are stored in a class variable and reset in the pytest fixture.
+    """
+
     requests = []
 
     def __init__(self, *args, **kwargs) -> None:
-        print("test")
         super().__init__(*args, **kwargs)
 
     def do_POST(self) -> None:
-        print(f"Received POST request: {self.path}")
         content_length = int(self.headers["Content-Length"])
         body = self.rfile.read(content_length).decode("utf-8")
-
         self.requests.append({"path": self.path, "headers": dict(self.headers), "body": body})
-
         self.send_response(200)
         self.end_headers()
         self.wfile.write(b"OK")
 
 
 @pytest.fixture
 def _llmobs_backend():
-    LLMObsBackend.requests = []
+    LLMObsServer.requests = []
     # Create and start the HTTP server
-    server = HTTPServer(("localhost", 0), LLMObsBackend)
+    server = HTTPServer(("localhost", 0), LLMObsServer)
     server_thread = threading.Thread(target=server.serve_forever)
     server_thread.daemon = True
     server_thread.start()
 
     # Provide the server details to the test
     server_address = f"http://{server.server_address[0]}:{server.server_address[1]}"
 
-    yield server_address, LLMObsBackend.requests
+    yield server_address, LLMObsServer.requests
 
     # Stop the server after the test
     server.shutdown()
     server.server_close()
 
 
 @pytest.fixture
-def llmobs_backend_requests(_llmobs_backend):
+def llmobs_backend(_llmobs_backend):
     _, reqs = _llmobs_backend
-    return reqs
+
+    class _LLMObsBackend:
+        def wait_for_num_events(self, num, attempts=1000):
+            for _ in range(attempts):
+                if len(reqs) == num:
+                    return [json.loads(r["body"]) for r in reqs]
+                # time.sleep will yield the GIL so the server can process the request
+                time.sleep(0.001)
+            else:
+                raise TimeoutError(f"Expected {num} events, got {len(reqs)}")
+
+    return _LLMObsBackend()
 
 
 @pytest.fixture

tests/llmobs/test_llmobs.py

@@ -247,17 +247,18 @@ def test_only_generate_span_events_from_llmobs_spans(tracer, llmobs_events):
     assert llmobs_events[0] == expected_grandchild_llmobs_span
 
 
-def test_utf_inputs_outputs(llmobs, llmobs_events, llmobs_backend_requests):
+def test_utf_inputs_outputs(llmobs, llmobs_backend):
     """Test that latin1 encoded inputs and outputs are correctly decoded."""
     with llmobs.llm(model_name="gpt-3.5-turbo-0125") as span:
         llmobs.annotate(
             span,
-            input_data="The first Super Bowl, which was formally known as the First AFL–NFL World Championship Game, was played on January 15, 1967.",
+            # uncomment to repro issue
+            # input_data="The first Super Bowl, which was formally known as the First AFL–NFL World Championship Game, was played on January 15, 1967.",
+            input_data="The first Super Bowl, which was formally known as the First AFL-NFL World Championship Game, was played on January 15, 1967.",
         )
 
-    import time
-
-    time.sleep(1)
-    assert len(llmobs_backend_requests) == 1
-    assert llmobs_backend_requests[0]
-    print(llmobs_backend_requests)
+    events = llmobs_backend.wait_for_num_events(num=1)
+    assert (
+        events[0]["spans"][0]["meta"]["input"]["messages"][0]["content"]
+        == "The first Super Bowl, which was formally known as the First AFL-NFL World Championship Game, was played on January 15, 1967."
+    )