mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 02:28:30 +08:00
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
import asyncio
|
|
import json
|
|
from collections.abc import AsyncGenerator
|
|
from typing import cast
|
|
|
|
from dify_agent.protocol.schemas import RunFailedEvent, RunFailedEventData, RunStartedEvent
|
|
from dify_agent.server.sse import format_sse_event, sse_event_stream
|
|
|
|
|
|
def test_format_sse_event_uses_id_event_and_json_data() -> None:
|
|
event = RunStartedEvent(id="7-0", run_id="run-1")
|
|
|
|
frame = format_sse_event(event)
|
|
|
|
assert frame.startswith("id: 7-0\nevent: run_started\ndata: ")
|
|
assert '"run_id":"run-1"' in frame
|
|
assert frame.endswith("\n\n")
|
|
|
|
|
|
def test_format_sse_event_escapes_unicode_line_separators() -> None:
|
|
error = "next-line:\x85line-separator:\u2028paragraph-separator:\u2029done"
|
|
event = RunFailedEvent(id="8-0", run_id="run-1", data=RunFailedEventData(error=error))
|
|
|
|
frame = format_sse_event(event)
|
|
data = next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: "))
|
|
|
|
assert "\x85" not in frame
|
|
assert "\u2028" not in frame
|
|
assert "\u2029" not in frame
|
|
assert "\\u0085" in frame
|
|
assert "\\u2028" in frame
|
|
assert "\\u2029" in frame
|
|
assert json.loads(data)["data"]["error"] == error
|
|
|
|
|
|
def test_sse_event_stream_emits_heartbeats_while_waiting() -> None:
|
|
async def scenario() -> None:
|
|
release = asyncio.Event()
|
|
|
|
async def events():
|
|
await release.wait()
|
|
yield RunStartedEvent(id="1-0", run_id="run-1")
|
|
|
|
stream = cast(AsyncGenerator[str, None], sse_event_stream(events(), heartbeat_interval_seconds=0.001))
|
|
assert await anext(stream) == ": keepalive\n\n"
|
|
|
|
_ = release.set()
|
|
assert (await anext(stream)).startswith("id: 1-0\nevent: run_started")
|
|
await stream.aclose()
|
|
|
|
asyncio.run(scenario())
|