117 lines
4.0 KiB
Python
117 lines
4.0 KiB
Python
import asyncio
|
|
import unittest
|
|
|
|
import zstandard
|
|
|
|
from sglang.srt.entrypoints.http_request_decompression import (
|
|
RequestDecompressionMiddleware,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
|
|
|
|
PAYLOAD = b'{"text":"hello world","n":7}'
|
|
COMPRESSED = zstandard.ZstdCompressor().compress(PAYLOAD)
|
|
|
|
|
|
def _drive(scope, body_chunks):
|
|
"""Drive the middleware once. Returns (seen, sent): `seen` is what the inner
|
|
app received ({scope, body}) or None if the app was never called; `sent` is
|
|
the list of ASGI messages the middleware emitted directly."""
|
|
seen = {}
|
|
sent = []
|
|
chunks = list(body_chunks)
|
|
|
|
async def receive():
|
|
if chunks:
|
|
chunk, more = chunks.pop(0)
|
|
return {"type": "http.request", "body": chunk, "more_body": more}
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message):
|
|
sent.append(message)
|
|
|
|
async def app(inner_scope, inner_receive, inner_send):
|
|
body = b""
|
|
more = True
|
|
while more:
|
|
message = await inner_receive()
|
|
if message["type"] != "http.request":
|
|
break
|
|
body += message.get("body", b"")
|
|
more = message.get("more_body", False)
|
|
seen["scope"] = inner_scope
|
|
seen["body"] = body
|
|
|
|
asyncio.run(RequestDecompressionMiddleware(app)(scope, receive, send))
|
|
return (seen or None), sent
|
|
|
|
|
|
class TestRequestDecompressionMiddleware(unittest.TestCase):
|
|
def test_passthrough_when_header_absent(self):
|
|
scope = {"type": "http", "headers": [(b"content-length", b"4")]}
|
|
seen, sent = _drive(scope, [(b"abcd", False)])
|
|
self.assertEqual(seen["body"], b"abcd")
|
|
self.assertEqual(seen["scope"]["headers"], [(b"content-length", b"4")])
|
|
self.assertEqual(sent, [])
|
|
|
|
def test_decompresses_zstd_body(self):
|
|
scope = {
|
|
"type": "http",
|
|
"headers": [
|
|
(b"x-body-compressed", b"zstd"),
|
|
(b"content-length", str(len(COMPRESSED)).encode()),
|
|
],
|
|
}
|
|
seen, sent = _drive(scope, [(COMPRESSED, False)])
|
|
self.assertEqual(seen["body"], PAYLOAD)
|
|
self.assertEqual(sent, [])
|
|
|
|
def test_strips_header_and_fixes_content_length(self):
|
|
scope = {
|
|
"type": "http",
|
|
"headers": [
|
|
(b"x-body-compressed", b"zstd"),
|
|
(b"content-length", str(len(COMPRESSED)).encode()),
|
|
(b"content-type", b"application/json"),
|
|
],
|
|
}
|
|
seen, _ = _drive(scope, [(COMPRESSED, False)])
|
|
self.assertEqual(
|
|
seen["scope"]["headers"],
|
|
[
|
|
(b"content-type", b"application/json"),
|
|
(b"content-length", str(len(PAYLOAD)).encode()),
|
|
],
|
|
)
|
|
|
|
def test_unsupported_method_returns_400(self):
|
|
scope = {"type": "http", "headers": [(b"x-body-compressed", b"gzip")]}
|
|
seen, sent = _drive(scope, [(b"abcd", False)])
|
|
self.assertIsNone(seen)
|
|
self.assertEqual(sent[0]["type"], "http.response.start")
|
|
self.assertEqual(sent[0]["status"], 400)
|
|
|
|
def test_chunked_body_reassembled(self):
|
|
half = len(COMPRESSED) // 2
|
|
scope = {"type": "http", "headers": [(b"x-body-compressed", b"zstd")]}
|
|
seen, _ = _drive(scope, [(COMPRESSED[:half], True), (COMPRESSED[half:], False)])
|
|
self.assertEqual(seen["body"], PAYLOAD)
|
|
|
|
def test_bad_body_returns_400(self):
|
|
scope = {"type": "http", "headers": [(b"x-body-compressed", b"zstd")]}
|
|
seen, sent = _drive(scope, [(b"not-zstd-data", False)])
|
|
self.assertIsNone(seen)
|
|
self.assertEqual(sent[0]["type"], "http.response.start")
|
|
self.assertEqual(sent[0]["status"], 400)
|
|
|
|
def test_non_http_scope_passthrough(self):
|
|
scope = {"type": "lifespan", "headers": []}
|
|
seen, sent = _drive(scope, [])
|
|
self.assertEqual(seen["scope"]["type"], "lifespan")
|
|
self.assertEqual(sent, [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|