79 lines
3.5 KiB
Python
79 lines
3.5 KiB
Python
"""Regression coverage for clients disconnecting while responses are sent."""
|
|
|
|
from email.message import Message
|
|
import errno
|
|
import io
|
|
from pathlib import Path
|
|
import time
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
from primethai.auth import AuthConfig
|
|
from primethai.webui import _handler_factory
|
|
|
|
|
|
class DashboardDisconnectTests(unittest.TestCase):
|
|
def make_handler(self, writer):
|
|
handler_type = _handler_factory(
|
|
(), Path("reports"), "2026-01-01", 10_000,
|
|
AuthConfig("test", "test", b"test-secret"),
|
|
)
|
|
handler = handler_type.__new__(handler_type)
|
|
handler.client_address = ("127.0.0.1", 12345)
|
|
handler.command = "GET"
|
|
handler.path = "/api/portfolio"
|
|
handler.requestline = "GET /api/portfolio HTTP/1.1"
|
|
handler.request_version = "HTTP/1.1"
|
|
handler.headers = Message()
|
|
handler.wfile = writer
|
|
handler.close_connection = False
|
|
handler._request_started_at = time.perf_counter()
|
|
handler._authenticated = Mock(return_value=True)
|
|
handler.log_message = Mock()
|
|
handler._request_log = Mock()
|
|
handler.send_response = Mock(wraps=handler.send_response)
|
|
return handler
|
|
|
|
def test_portfolio_disconnect_does_not_attempt_a_second_response(self):
|
|
for error_type in (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
|
|
for failing_write in (1, 2): # Header flush or response body.
|
|
with self.subTest(error=error_type, failing_write=failing_write):
|
|
writer = Mock()
|
|
writer.write.side_effect = (
|
|
[error_type(errno.EPIPE, "connection closed")]
|
|
if failing_write == 1
|
|
else [100, error_type(errno.ECONNRESET, "connection closed")]
|
|
)
|
|
handler = self.make_handler(writer)
|
|
with patch("primethai.webui._portfolio_data", return_value={"cash": 100}):
|
|
handler.do_GET()
|
|
handler.send_response.assert_called_once_with(200)
|
|
self.assertTrue(handler.close_connection)
|
|
self.assertEqual(handler._request_log.call_args.args, ("client_disconnected",))
|
|
|
|
def test_redirect_handles_closed_connection(self):
|
|
writer = Mock()
|
|
writer.write.side_effect = BrokenPipeError(errno.EPIPE, "connection closed")
|
|
handler = self.make_handler(writer)
|
|
handler._redirect("/login")
|
|
handler.send_response.assert_called_once_with(303)
|
|
self.assertTrue(handler.close_connection)
|
|
|
|
def test_success_and_actual_portfolio_failure_remain_distinct(self):
|
|
for failed in (False, True):
|
|
with self.subTest(failed=failed):
|
|
handler = self.make_handler(io.BytesIO())
|
|
with patch("primethai.webui._portfolio_data", return_value={"cash": 100},
|
|
side_effect=ValueError("invalid portfolio") if failed else None):
|
|
handler.do_GET()
|
|
handler.send_response.assert_called_once_with(500 if failed else 200)
|
|
response = handler.wfile.getvalue()
|
|
self.assertIn(b"invalid portfolio" if failed else b'"cash": 100', response)
|
|
self.assertFalse(handler.close_connection)
|
|
if failed:
|
|
self.assertEqual(handler._request_log.call_args.args, ("portfolio_failed",))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|