78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import asyncio
|
|
|
|
from app.models.health import HealthSnapshot
|
|
from app.models.proxy import ProxyDocument
|
|
from app.models.service import ServiceEndpointDocument
|
|
from app.services.mitmproxy.health import check_proxy
|
|
|
|
|
|
def test_check_proxy_reports_upstream_and_proxy_healthy():
|
|
async def run():
|
|
upstream = await asyncio.start_server(http_handler, "127.0.0.1", 0)
|
|
proxy = await asyncio.start_server(http_handler, "127.0.0.1", 0)
|
|
try:
|
|
upstream_port = upstream.sockets[0].getsockname()[1]
|
|
proxy_port = proxy.sockets[0].getsockname()[1]
|
|
document = ProxyDocument(
|
|
proxy_port=proxy_port,
|
|
name="health",
|
|
code="",
|
|
service=ServiceEndpointDocument(port=upstream_port),
|
|
)
|
|
|
|
health = await check_proxy(document)
|
|
finally:
|
|
upstream.close()
|
|
proxy.close()
|
|
await upstream.wait_closed()
|
|
await proxy.wait_closed()
|
|
|
|
assert isinstance(health, HealthSnapshot)
|
|
assert health.overall == "healthy"
|
|
assert [check.target for check in health.checks] == [
|
|
"upstream",
|
|
"proxy",
|
|
]
|
|
assert all(check.duration_ms is None for check in health.checks)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_check_proxy_reports_benchmark_latency_when_enabled():
|
|
async def run():
|
|
upstream = await asyncio.start_server(http_handler, "127.0.0.1", 0)
|
|
proxy = await asyncio.start_server(http_handler, "127.0.0.1", 0)
|
|
try:
|
|
upstream_port = upstream.sockets[0].getsockname()[1]
|
|
proxy_port = proxy.sockets[0].getsockname()[1]
|
|
document = ProxyDocument(
|
|
proxy_port=proxy_port,
|
|
name="health",
|
|
code="",
|
|
healthcheck={"benchmark_enabled": True},
|
|
service=ServiceEndpointDocument(port=upstream_port),
|
|
)
|
|
|
|
health = await check_proxy(document)
|
|
finally:
|
|
upstream.close()
|
|
proxy.close()
|
|
await upstream.wait_closed()
|
|
await proxy.wait_closed()
|
|
|
|
assert isinstance(health, HealthSnapshot)
|
|
assert all(
|
|
check.duration_ms is not None and check.duration_ms >= 0
|
|
for check in health.checks
|
|
)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
async def http_handler(reader, writer):
|
|
await reader.readline()
|
|
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")
|
|
await writer.drain()
|
|
writer.close()
|
|
await writer.wait_closed()
|