Add a test for eventstreams

Author: yuvipanda

tests/resources/eventstream.py

@@ -0,0 +1,36 @@
+import asyncio
+
+import tornado.escape
+import tornado.ioloop
+import tornado.options
+import tornado.web
+import tornado.websocket
+from tornado.options import define, options
+
+
+class Application(tornado.web.Application):
+    def __init__(self):
+        handlers = [
+            (r"/stream/(\d+)", StreamHandler),
+        ]
+        super().__init__(handlers)
+
+
+class StreamHandler(tornado.web.RequestHandler):
+    async def get(self, seconds):
+        for i in range(int(seconds)):
+            await asyncio.sleep(0.5)
+            self.write(f"data: {i}\n\n")
+            await self.flush()
+
+
+def main():
+    define("port", default=8888, help="run on the given port", type=int)
+    options.parse_command_line()
+    app = Application()
+    app.listen(options.port)
+    tornado.ioloop.IOLoop.current().start()
+
+
+if __name__ == "__main__":
+    main()

tests/resources/jupyter_server_config.py

@@ -79,6 +79,9 @@ def my_env():
             "X-Custom-Header": "pytest-23456",
         },
     },
+    "python-eventstream": {
+        "command": [sys.executable, "./tests/resources/eventstream.py", "--port={port}"]
+    },
     "python-unix-socket-true": {
         "command": [
             sys.executable,

tests/test_proxies.py

@@ -1,13 +1,14 @@
 import gzip
 import json
 import sys
+import time
 from http.client import HTTPConnection
 from io import BytesIO
 from typing import Tuple
 from urllib.parse import quote
 
 import pytest
-from tornado.httpclient import HTTPClientError
+from tornado.httpclient import AsyncHTTPClient, HTTPClientError
 from tornado.websocket import websocket_connect
 
 # use ipv4 for CI, etc.
@@ -343,6 +344,39 @@ def test_server_content_encoding_header(
         assert f.read() == b"this is a test"
 
 
+async def test_eventstream(a_server_port_and_token: Tuple[int, str]) -> None:
+    PORT, TOKEN = a_server_port_and_token
+    # The test server under eventstream.py will send back monotonically increasing numbers
+    # starting at 0 until the specified limit, with a 500ms gap between them. We test that:
+    # 1. We get back as many callbacks from our streaming read as the total number,
+    #    as the server does a flush after writing each entry.
+    # 2. The streaming entries are read (with some error margin) around the 500ms mark, to
+    #    ensure this is *actually* being streamed
+    limit = 3
+    last_cb_time = time.perf_counter()
+    times_called = 0
+    stream_read_intervals = []
+
+    def streaming_cb(data):
+        nonlocal times_called, last_cb_time, stream_read_intervals
+        time_taken = time.perf_counter() - last_cb_time
+        last_cb_time = time.perf_counter()
+        stream_read_intervals.append(time_taken)
+        times_called += 1
+
+    url = f"http://{LOCALHOST}:{PORT}/python-eventstream/stream/{limit}?token={TOKEN}"
+    client = AsyncHTTPClient()
+    await client.fetch(
+        url,
+        headers={"Accept": "text/event-stream"},
+        request_timeout=22,
+        streaming_callback=streaming_cb,
+    )
+    assert times_called == limit
+    print(stream_read_intervals)
+    assert all([0.45 < t < 0.7 for t in stream_read_intervals])
+
+
 async def test_server_proxy_websocket_messages(
     a_server_port_and_token: Tuple[int, str]
 ) -> None: