|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +Test NDJSON streaming response (server -> client) |
| 4 | +""" |
| 5 | +import unittest |
| 6 | +import socket |
| 7 | +import time |
| 8 | +import json |
| 9 | +import threading |
| 10 | +from uhttp import server as uhttp_server |
| 11 | + |
| 12 | + |
| 13 | +class TestNDJSON(unittest.TestCase): |
| 14 | + """Test suite for NDJSON streaming responses""" |
| 15 | + |
| 16 | + server = None |
| 17 | + server_thread = None |
| 18 | + nd_clients = [] |
| 19 | + PORT = 9986 |
| 20 | + |
| 21 | + @classmethod |
| 22 | + def setUpClass(cls): |
| 23 | + cls.server = uhttp_server.HttpServer(port=cls.PORT) |
| 24 | + |
| 25 | + def run_server(): |
| 26 | + try: |
| 27 | + while cls.server: |
| 28 | + client = cls.server.wait(timeout=0.1) |
| 29 | + |
| 30 | + if client: |
| 31 | + if client.path == '/stream': |
| 32 | + if client.response_ndjson(): |
| 33 | + cls.nd_clients.append({ |
| 34 | + 'client': client, |
| 35 | + 'counter': 0, |
| 36 | + 'last_send': time.time(), |
| 37 | + 'mode': 'dict', |
| 38 | + }) |
| 39 | + elif client.path == '/stream-mixed': |
| 40 | + if client.response_ndjson(): |
| 41 | + cls.nd_clients.append({ |
| 42 | + 'client': client, |
| 43 | + 'counter': 0, |
| 44 | + 'last_send': time.time(), |
| 45 | + 'mode': 'mixed', |
| 46 | + }) |
| 47 | + elif client.path == '/stream-headers': |
| 48 | + if client.response_ndjson( |
| 49 | + headers={'X-Stream': 'ndjson'}): |
| 50 | + cls.nd_clients.append({ |
| 51 | + 'client': client, |
| 52 | + 'counter': 0, |
| 53 | + 'last_send': time.time(), |
| 54 | + 'mode': 'dict', |
| 55 | + }) |
| 56 | + else: |
| 57 | + client.respond("Not found", status=404) |
| 58 | + |
| 59 | + for sc in list(cls.nd_clients): |
| 60 | + if time.time() - sc['last_send'] > 0.05: |
| 61 | + sc['counter'] += 1 |
| 62 | + mode = sc['mode'] |
| 63 | + |
| 64 | + if sc['counter'] >= 4: |
| 65 | + sc['client'].response_stream_end() |
| 66 | + cls.nd_clients.remove(sc) |
| 67 | + elif mode == 'dict': |
| 68 | + sc['client'].send_ndjson( |
| 69 | + {'n': sc['counter'], 'msg': 'hello'}) |
| 70 | + sc['last_send'] = time.time() |
| 71 | + elif mode == 'mixed': |
| 72 | + # rotate through different JSON types |
| 73 | + values = [ |
| 74 | + {'n': sc['counter']}, |
| 75 | + [1, 2, sc['counter']], |
| 76 | + f'string {sc["counter"]}', |
| 77 | + sc['counter']] |
| 78 | + sc['client'].send_ndjson( |
| 79 | + values[(sc['counter'] - 1) % len(values)]) |
| 80 | + sc['last_send'] = time.time() |
| 81 | + |
| 82 | + except Exception: |
| 83 | + pass |
| 84 | + |
| 85 | + cls.server_thread = threading.Thread(target=run_server, daemon=True) |
| 86 | + cls.server_thread.start() |
| 87 | + time.sleep(0.5) |
| 88 | + |
| 89 | + @classmethod |
| 90 | + def tearDownClass(cls): |
| 91 | + if cls.server: |
| 92 | + cls.server.close() |
| 93 | + cls.server = None |
| 94 | + |
| 95 | + def setUp(self): |
| 96 | + TestNDJSON.nd_clients = [] |
| 97 | + |
| 98 | + def _recv_all(self, sock, timeout=3.0): |
| 99 | + sock.settimeout(timeout) |
| 100 | + all_data = b"" |
| 101 | + try: |
| 102 | + while True: |
| 103 | + chunk = sock.recv(4096) |
| 104 | + if not chunk: |
| 105 | + break |
| 106 | + all_data += chunk |
| 107 | + except socket.timeout: |
| 108 | + pass |
| 109 | + return all_data |
| 110 | + |
| 111 | + def _make_request(self, path): |
| 112 | + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 113 | + sock.connect(('localhost', self.PORT)) |
| 114 | + request = ( |
| 115 | + f"GET {path} HTTP/1.1\r\n" |
| 116 | + f"Host: localhost\r\n" |
| 117 | + f"Connection: close\r\n" |
| 118 | + f"\r\n" |
| 119 | + ).encode() |
| 120 | + sock.sendall(request) |
| 121 | + return sock |
| 122 | + |
| 123 | + def _split_response(self, data): |
| 124 | + """Split raw HTTP response into (headers_text, body_bytes)""" |
| 125 | + sep = data.find(b"\r\n\r\n") |
| 126 | + self.assertGreater(sep, 0, "no header/body separator") |
| 127 | + return data[:sep].decode(), data[sep + 4:] |
| 128 | + |
| 129 | + def test_stream_headers(self): |
| 130 | + """NDJSON response has correct content-type and cache-control""" |
| 131 | + sock = self._make_request('/stream') |
| 132 | + try: |
| 133 | + data = self._recv_all(sock) |
| 134 | + headers, _ = self._split_response(data) |
| 135 | + self.assertIn("200 OK", headers) |
| 136 | + self.assertIn("content-type: application/x-ndjson", headers) |
| 137 | + self.assertIn("cache-control: no-cache", headers) |
| 138 | + finally: |
| 139 | + sock.close() |
| 140 | + |
| 141 | + def test_custom_headers_passthrough(self): |
| 142 | + """Custom headers passed to response_ndjson appear in response""" |
| 143 | + sock = self._make_request('/stream-headers') |
| 144 | + try: |
| 145 | + data = self._recv_all(sock) |
| 146 | + headers, _ = self._split_response(data) |
| 147 | + self.assertIn("X-Stream: ndjson", headers) |
| 148 | + self.assertIn("content-type: application/x-ndjson", headers) |
| 149 | + finally: |
| 150 | + sock.close() |
| 151 | + |
| 152 | + def test_ndjson_lines_parse(self): |
| 153 | + """Each body line is a valid JSON object terminated by \\n""" |
| 154 | + sock = self._make_request('/stream') |
| 155 | + try: |
| 156 | + data = self._recv_all(sock) |
| 157 | + _, body = self._split_response(data) |
| 158 | + |
| 159 | + # body must end with \n on the last record |
| 160 | + self.assertTrue(body.endswith(b"\n")) |
| 161 | + |
| 162 | + lines = body.split(b"\n") |
| 163 | + # last element is empty string after trailing \n |
| 164 | + self.assertEqual(lines[-1], b"") |
| 165 | + records = [json.loads(l) for l in lines[:-1]] |
| 166 | + |
| 167 | + self.assertEqual(len(records), 3) |
| 168 | + for i, rec in enumerate(records, start=1): |
| 169 | + self.assertEqual(rec, {'n': i, 'msg': 'hello'}) |
| 170 | + finally: |
| 171 | + sock.close() |
| 172 | + |
| 173 | + def test_no_embedded_newlines(self): |
| 174 | + """Each NDJSON line contains exactly one record (no embedded \\n)""" |
| 175 | + sock = self._make_request('/stream') |
| 176 | + try: |
| 177 | + data = self._recv_all(sock) |
| 178 | + _, body = self._split_response(data) |
| 179 | + |
| 180 | + for line in body.split(b"\n")[:-1]: |
| 181 | + self.assertNotIn(b"\n", line) |
| 182 | + # must be parseable on its own |
| 183 | + json.loads(line) |
| 184 | + finally: |
| 185 | + sock.close() |
| 186 | + |
| 187 | + def test_mixed_json_types(self): |
| 188 | + """send_ndjson accepts dict/list/str/int per-record""" |
| 189 | + sock = self._make_request('/stream-mixed') |
| 190 | + try: |
| 191 | + data = self._recv_all(sock) |
| 192 | + _, body = self._split_response(data) |
| 193 | + lines = body.split(b"\n")[:-1] |
| 194 | + self.assertEqual(len(lines), 3) |
| 195 | + self.assertEqual(json.loads(lines[0]), {'n': 1}) |
| 196 | + self.assertEqual(json.loads(lines[1]), [1, 2, 2]) |
| 197 | + self.assertEqual(json.loads(lines[2]), 'string 3') |
| 198 | + finally: |
| 199 | + sock.close() |
| 200 | + |
| 201 | + def test_stream_end_closes_connection(self): |
| 202 | + """response_stream_end closes the socket after final record""" |
| 203 | + sock = self._make_request('/stream') |
| 204 | + try: |
| 205 | + data = self._recv_all(sock, timeout=3.0) |
| 206 | + _, body = self._split_response(data) |
| 207 | + self.assertTrue(len(body) > 0) |
| 208 | + remaining = sock.recv(1024) |
| 209 | + self.assertEqual(remaining, b"") |
| 210 | + finally: |
| 211 | + sock.close() |
| 212 | + |
| 213 | + |
| 214 | +if __name__ == '__main__': |
| 215 | + unittest.main() |
0 commit comments