135 lines
4.6 KiB
Python
135 lines
4.6 KiB
Python
"""DIP-003 keep-alive ownership and framing-adjacent transport tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
|
|
import pytest
|
|
|
|
|
|
class _Response:
|
|
def __init__(self, status: int = 200, body: bytes = b"ok", headers: dict | None = None):
|
|
self.status = status
|
|
self._body = body
|
|
self.headers = headers or {"Content-Type": "application/octet-stream", "Content-Length": str(len(body))}
|
|
|
|
def read(self) -> bytes:
|
|
return self._body
|
|
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
|
|
class _Connection:
|
|
instances: list["_Connection"] = []
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.requests: list[tuple] = []
|
|
self.responses: list[object] = [_Response(), _Response()]
|
|
self.closed = False
|
|
self.__class__.instances.append(self)
|
|
|
|
def request(self, *args, **kwargs) -> None:
|
|
self.requests.append((args, kwargs))
|
|
|
|
def getresponse(self):
|
|
response = self.responses.pop(0)
|
|
if isinstance(response, BaseException):
|
|
raise response
|
|
return response
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
def test_direct_hop_client_reuses_one_connection_and_discards_uncertain_socket(monkeypatch):
|
|
"""A Route Session owns one serialized direct socket and never replays it.
|
|
|
|
Tags: performance, routing
|
|
"""
|
|
from meshnet_node import torch_server
|
|
|
|
_Connection.instances = []
|
|
monkeypatch.setattr(torch_server.http.client, "HTTPConnection", _Connection)
|
|
client = torch_server._DirectHopClient("http://tail.example:8001")
|
|
|
|
assert client.request("/forward", b"one", {})[2] == b"ok"
|
|
assert client.request("/forward", b"two", {})[2] == b"ok"
|
|
assert len(_Connection.instances) == 1
|
|
assert [call[0][1] for call in _Connection.instances[0].requests] == ["/forward", "/forward"]
|
|
|
|
_Connection.instances[0].responses.append(ConnectionResetError("stale peer"))
|
|
with pytest.raises(torch_server._DirectRequestUncertainError):
|
|
client.request("/forward", b"three", {})
|
|
assert _Connection.instances[0].closed is True
|
|
assert client._connection is None
|
|
|
|
|
|
def test_bridge_pool_reuses_a_worker_connection_and_invalidates_stale_one(monkeypatch):
|
|
"""A bridge worker keeps its own loopback client; broken clients are dropped.
|
|
|
|
Tags: performance, relay
|
|
"""
|
|
from meshnet_node import relay_bridge
|
|
|
|
_Connection.instances = []
|
|
monkeypatch.setattr(relay_bridge.http.client, "HTTPConnection", _Connection)
|
|
bridge = relay_bridge.RelayHttpBridge(
|
|
relay_url="ws://relay.example/ws",
|
|
peer_id="peer",
|
|
local_base_url="http://127.0.0.1:8001",
|
|
advertised_addr="",
|
|
)
|
|
frames: list[dict] = []
|
|
bridge._send_response_frame = lambda frame: (frames.append(frame), True)[1]
|
|
|
|
request = {"request_id": "one", "method": "POST", "path": "/forward", "headers": {}, "body": ""}
|
|
bridge._process_request(request)
|
|
bridge._process_request({**request, "request_id": "two"})
|
|
assert len(_Connection.instances) == 1
|
|
assert len(_Connection.instances[0].requests) == 2
|
|
|
|
_Connection.instances[0].responses.append(ConnectionResetError("stale loopback"))
|
|
bridge._process_request({**request, "request_id": "broken"})
|
|
assert _Connection.instances[0].closed is True
|
|
bridge._process_request({**request, "request_id": "replacement"})
|
|
assert len(_Connection.instances) == 2
|
|
bridge.stop()
|
|
|
|
|
|
def test_node_sse_uses_chunked_framing_and_tolerates_client_cancellation():
|
|
"""HTTP/1.1 streams terminate without EOF and ignore a cancelled client.
|
|
|
|
Tags: performance, streaming
|
|
"""
|
|
from meshnet_node.torch_server import _TorchHandler
|
|
|
|
handler = object.__new__(_TorchHandler)
|
|
headers: list[tuple[str, str]] = []
|
|
handler.send_response = lambda status: None
|
|
handler.send_header = lambda name, value: headers.append((name, value))
|
|
handler.end_headers = lambda: None
|
|
handler.wfile = io.BytesIO()
|
|
emit = handler._start_openai_stream("model")
|
|
emit(None)
|
|
wire = handler.wfile.getvalue()
|
|
|
|
assert _TorchHandler.protocol_version == "HTTP/1.1"
|
|
assert ("Transfer-Encoding", "chunked") in headers
|
|
assert wire.endswith(b"0\r\n\r\n")
|
|
assert b"data: [DONE]" in wire
|
|
|
|
class _CancelledWriter:
|
|
def write(self, _body):
|
|
raise BrokenPipeError
|
|
|
|
def flush(self):
|
|
raise BrokenPipeError
|
|
|
|
cancelled = object.__new__(_TorchHandler)
|
|
cancelled.send_response = lambda status: None
|
|
cancelled.send_header = lambda name, value: None
|
|
cancelled.end_headers = lambda: None
|
|
cancelled.wfile = _CancelledWriter()
|
|
cancelled._start_openai_stream("model")(None)
|