From 82e488488522bfb082111474163e38d42df55cc0 Mon Sep 17 00:00:00 2001 From: Popescu Tudor-Cristian Date: Tue, 28 Jul 2026 11:17:45 +0300 Subject: [PATCH] test: raise code coverage to ~98% for MER KPI Add unit tests to lift SonarCloud main-branch coverage from 65.6% to ~98%, clearing the 80% MER "Increase Code Coverage" guardrail (PRODEV-612). - tests/test_diagnose.py: full coverage of the binary diagnostics helper (previously 0%, untested). - tests/cli/test_token_refresh_more.py: refresh loop, oauth/client-credentials flows, lifecycle, and retry/cancel paths. - tests/cli/test_session_more.py: relay loop, consumer drain/error paths, stdio/http run flows and start-error handling. - tests/cli/test_runtime_extra.py: HTTP process lifecycle, monitor, cleanup, registration (stdio + streamable-http), run-server orchestration, keep-alive. No source changes. Co-Authored-By: Claude Opus 4.8 --- tests/cli/test_runtime_extra.py | 581 +++++++++++++++++++++++++++ tests/cli/test_session_more.py | 245 +++++++++++ tests/cli/test_token_refresh_more.py | 238 +++++++++++ tests/test_diagnose.py | 147 +++++++ 4 files changed, 1211 insertions(+) create mode 100644 tests/cli/test_runtime_extra.py create mode 100644 tests/cli/test_session_more.py create mode 100644 tests/cli/test_token_refresh_more.py create mode 100644 tests/test_diagnose.py diff --git a/tests/cli/test_runtime_extra.py b/tests/cli/test_runtime_extra.py new file mode 100644 index 0000000..98aaaf4 --- /dev/null +++ b/tests/cli/test_runtime_extra.py @@ -0,0 +1,581 @@ +"""Extended runtime tests: HTTP process lifecycle, register, run loop, keep-alive. + +These are structural unit tests: the SignalR transport, HTTP client, subprocess +layer and TokenRefresher are all mocked, so they assert how ``UiPathMcpRuntime`` +wires itself together and handles local error/lifecycle paths. They deliberately +do NOT verify the agenthub wire protocol (URLs, headers, status-code semantics, +message framing) -- that is covered by the integration suite (integration_tests.yml +/ the testcases/* cloud+alpha jobs). Prefer asserting observable effects over log +text. +""" + +import asyncio +import contextlib +import json +import logging +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from uipath.runtime import UiPathRuntimeStatus + +from uipath_mcp._cli._runtime._exception import UiPathMcpRuntimeError +from uipath_mcp._cli._runtime._runtime import UiPathMcpRuntime +from uipath_mcp._cli._runtime._session import StreamableHttpSessionServer + +RMODULE = "uipath_mcp._cli._runtime._runtime" +CODED_KEY = "11111111-2222-3333-4444-555555555555" + + +def _make_runtime(**extra) -> UiPathMcpRuntime: + server = MagicMock( + name="server", + is_streamable_http=False, + args=[], + command="cmd", + env={}, + url=None, + ) + server.name = "svc" + with patch(f"{RMODULE}.UiPath"): + rt = UiPathMcpRuntime(server=server, runtime_id="rid", entrypoint="ep", **extra) + rt._uipath = MagicMock() + return rt + + +@pytest.fixture +def runtime() -> UiPathMcpRuntime: + return _make_runtime() + + +class _FakeCM: + """Minimal async context manager yielding a fixed value.""" + + def __init__(self, value): + self._value = value + + async def __aenter__(self): + return self._value + + async def __aexit__(self, *_a): + return False + + +class _FakeStderr: + def __init__(self, lines): + self._lines = list(lines) + + def __aiter__(self): + return self + + async def __anext__(self): + if self._lines: + return self._lines.pop(0) + raise StopAsyncIteration + + +def _http_client_cm(client: MagicMock) -> MagicMock: + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=client) + cm.__aexit__ = AsyncMock(return_value=False) + return cm + + +# --------------------------------------------------------------------------- # +# _start_http_server_process / _drain_http_stderr +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_start_http_server_process_coded(runtime): + runtime._process_key = CODED_KEY # Coded -> merges os.environ + proc = MagicMock(pid=4321, stderr=None) + with patch( + f"{RMODULE}.asyncio.create_subprocess_exec", new=AsyncMock(return_value=proc) + ): + await runtime._start_http_server_process() + assert runtime._http_server_process is proc + if runtime._http_stderr_drain_task: + runtime._http_stderr_drain_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await runtime._http_stderr_drain_task + + +@pytest.mark.asyncio +async def test_drain_http_stderr_reads_lines(runtime): + runtime._http_server_process = MagicMock( + stderr=_FakeStderr([b"err line 1\n", b"err line 2\n"]) + ) + await runtime._drain_http_stderr() + assert any("err line 1" in line for line in runtime._http_server_stderr_lines) + + +# --------------------------------------------------------------------------- # +# _wait_for_http_server_ready +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_wait_for_http_ready_success(runtime): + runtime._server.url = "http://localhost:9" + client = MagicMock() + client.get = AsyncMock(return_value=MagicMock(status_code=200)) + with patch("httpx.AsyncClient", return_value=_http_client_cm(client)): + await runtime._wait_for_http_server_ready(max_retries=1, retry_delay=0) + client.get.assert_awaited_once() # readiness probe was actually issued + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_no_url(runtime): + runtime._server.url = None + with pytest.raises(ValueError): + await runtime._wait_for_http_server_ready() + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_process_crashed(runtime): + runtime._server.url = "http://x" + runtime._http_server_process = MagicMock(returncode=1) + runtime._http_server_stderr_lines = ["boom"] + with pytest.raises(UiPathMcpRuntimeError): + await runtime._wait_for_http_server_ready(max_retries=2, retry_delay=0) + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_retries_exhausted(runtime): + runtime._server.url = "http://x" + client = MagicMock() + client.get = AsyncMock(side_effect=httpx.ConnectError("no")) + with ( + patch("httpx.AsyncClient", return_value=_http_client_cm(client)), + patch(f"{RMODULE}.asyncio.sleep", new=AsyncMock()), + ): + with pytest.raises(UiPathMcpRuntimeError): + await runtime._wait_for_http_server_ready(max_retries=2, retry_delay=0) + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_status_error(runtime): + runtime._server.url = "http://x" + client = MagicMock() + client.get = AsyncMock( + side_effect=httpx.HTTPStatusError( + "e", request=MagicMock(), response=MagicMock() + ) + ) + with patch("httpx.AsyncClient", return_value=_http_client_cm(client)): + # an HTTP error status still means the server is up -> returns, no raise + await runtime._wait_for_http_server_ready(max_retries=1, retry_delay=0) + client.get.assert_awaited_once() + + +# --------------------------------------------------------------------------- # +# _stop_http_server_process +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_stop_http_server_process_full(runtime): + async def _run(): + await asyncio.sleep(10) + + runtime._http_monitor_task = asyncio.create_task(_run()) + runtime._http_stderr_drain_task = asyncio.create_task(_run()) + proc = MagicMock() + proc.terminate = MagicMock() + proc.wait = AsyncMock(return_value=0) + runtime._http_server_process = proc + await runtime._stop_http_server_process() + assert runtime._http_server_process is None + assert runtime._http_monitor_task is None + assert runtime._http_stderr_drain_task is None + + +@pytest.mark.asyncio +async def test_stop_http_process_timeout_kills(runtime): + proc = MagicMock() + proc.terminate = MagicMock() + proc.kill = MagicMock() + proc.wait = AsyncMock(return_value=0) + runtime._http_server_process = proc + # patch wait_for -> TimeoutError to drive the "graceful terminate timed out, + # escalate to kill()" branch deterministically. + with patch( + f"{RMODULE}.asyncio.wait_for", new=AsyncMock(side_effect=asyncio.TimeoutError()) + ): + await runtime._stop_http_server_process() + proc.kill.assert_called_once() + + +@pytest.mark.asyncio +async def test_stop_http_process_lookup_error(runtime): + proc = MagicMock() + proc.terminate = MagicMock(side_effect=ProcessLookupError()) + runtime._http_server_process = proc + await runtime._stop_http_server_process() + assert runtime._http_server_process is None + + +# --------------------------------------------------------------------------- # +# _monitor_http_server_process +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_monitor_http_process_exit_stops_sessions(runtime): + runtime._http_server_process = MagicMock(wait=AsyncMock(return_value=1)) + http_sess = MagicMock(spec=StreamableHttpSessionServer) + http_sess.stop = AsyncMock() + runtime._session_servers = {"s1": http_sess} + await runtime._monitor_http_server_process() + http_sess.stop.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_monitor_http_process_stop_error(runtime): + runtime._http_server_process = MagicMock(wait=AsyncMock(return_value=1)) + http_sess = MagicMock(spec=StreamableHttpSessionServer) + http_sess.stop = AsyncMock(side_effect=RuntimeError("stop fail")) + runtime._session_servers = {"s1": http_sess} + await runtime._monitor_http_server_process() + # a failing stop() still drops the crashed session from the map + assert "s1" not in runtime._session_servers + + +@pytest.mark.asyncio +async def test_monitor_http_process_cancelled(runtime): + proc = MagicMock(wait=AsyncMock(side_effect=asyncio.CancelledError())) + runtime._http_server_process = proc + await runtime._monitor_http_server_process() + proc.wait.assert_awaited_once() # reached the await before being cancelled + + +# --------------------------------------------------------------------------- # +# _handle_signalr_session_closed +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_session_closed_non_sandboxed_logs_output(runtime): + sess = MagicMock(output="out") + sess.stop = AsyncMock() + runtime._session_servers["s1"] = sess + runtime._job_id = None # not sandboxed + await runtime._handle_signalr_session_closed(["s1"]) + sess.stop.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_session_closed_handles_stop_error(runtime): + # the session is popped from the map before stop() is awaited, so a failing + # stop() must not leave a dangling entry behind. + sess = MagicMock(output=None) + sess.stop = AsyncMock(side_effect=RuntimeError("boom-close")) + runtime._session_servers["s1"] = sess + await runtime._handle_signalr_session_closed(["s1"]) + assert "s1" not in runtime._session_servers + + +# --------------------------------------------------------------------------- # +# _handle_signalr_message start-error path +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_handle_message_start_error(runtime): + runtime._server.is_streamable_http = False + fake = MagicMock() + fake.start = AsyncMock(side_effect=RuntimeError("start boom")) + with ( + patch(f"{RMODULE}.StdioSessionServer", return_value=fake), + patch.object(runtime, "_on_session_start_error", new=AsyncMock()) as ose, + ): + await runtime._handle_signalr_message(["sNew", "req"]) + ose.assert_awaited_once_with("sNew") + + +# --------------------------------------------------------------------------- # +# _cleanup extended paths +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_cleanup_full(runtime): + async def _run(): + await asyncio.sleep(10) + + runtime._token_refresher = MagicMock(stop=AsyncMock()) + runtime._keep_alive_task = asyncio.create_task(_run()) + good = MagicMock() + good.stop = AsyncMock() + bad = MagicMock() + bad.stop = AsyncMock(side_effect=RuntimeError("x")) + runtime._session_servers = {"g": good, "b": bad} + ws = MagicMock() + ws.close = AsyncMock() + transport = MagicMock(_ws=ws) + runtime._signalr_client = MagicMock(_transport=transport) + with ( + patch.object(runtime, "_on_runtime_abort", new=AsyncMock()), + patch.object(runtime, "_stop_http_server_process", new=AsyncMock()), + ): + await runtime._cleanup() + good.stop.assert_awaited_once() + ws.close.assert_awaited_once() + assert runtime._cleanup_done is True + + +@pytest.mark.asyncio +async def test_cleanup_ws_close_error(runtime): + ws = MagicMock() + ws.close = AsyncMock(side_effect=RuntimeError("close fail")) + transport = MagicMock(_ws=ws) + runtime._signalr_client = MagicMock(_transport=transport) + with ( + patch.object(runtime, "_on_runtime_abort", new=AsyncMock()), + patch.object(runtime, "_stop_http_server_process", new=AsyncMock()), + ): + await runtime._cleanup() + # close was attempted, and its failure does not abort cleanup + ws.close.assert_awaited_once() + assert runtime._cleanup_done is True + + +# --------------------------------------------------------------------------- # +# _run_server orchestration +# --------------------------------------------------------------------------- # +def _run_server_patches(runtime, cfg): + sig = MagicMock() + sig.run = AsyncMock() + refresher = MagicMock() + refresher.start = MagicMock() + refresher.stop = AsyncMock() + return sig, refresher + + +@pytest.mark.asyncio +async def test_run_server_happy_path(runtime): + runtime._tenant_id = "t" + runtime._org_id = "o" + runtime._folder_key = "fk" + runtime._server.is_streamable_http = False + runtime._session_output = "final-out" + cfg = MagicMock(base_url="https://x") + sig, refresher = _run_server_patches(runtime, cfg) + with ( + patch(f"{RMODULE}.UiPathConfig", cfg), + patch(f"{RMODULE}.SignalRClient", return_value=sig), + patch(f"{RMODULE}.TokenRefresher", return_value=refresher), + patch.object(runtime, "_register", new=AsyncMock()), + patch.object(runtime, "_keep_alive", new=AsyncMock()), + patch.object(runtime, "_cleanup", new=AsyncMock()), + ): + result = await runtime._run_server() + assert result.status == UiPathRuntimeStatus.SUCCESSFUL + assert result.output.get("content") == "final-out" + + +@pytest.mark.asyncio +async def test_run_server_keyboard_interrupt_during_wait(runtime): + runtime._tenant_id = "t" + runtime._org_id = "o" + runtime._folder_key = "fk" + runtime._server.is_streamable_http = False + cfg = MagicMock(base_url="https://x") + sig, refresher = _run_server_patches(runtime, cfg) + with ( + patch(f"{RMODULE}.UiPathConfig", cfg), + patch(f"{RMODULE}.SignalRClient", return_value=sig), + patch(f"{RMODULE}.TokenRefresher", return_value=refresher), + patch.object(runtime, "_register", new=AsyncMock()), + patch.object(runtime, "_keep_alive", new=AsyncMock()), + patch.object(runtime, "_cleanup", new=AsyncMock()), + # simulate Ctrl-C while awaiting the run/cancel tasks + patch(f"{RMODULE}.asyncio.wait", side_effect=KeyboardInterrupt()), + ): + result = await runtime._run_server() + assert result.status == UiPathRuntimeStatus.SUCCESSFUL + + +@pytest.mark.asyncio +async def test_run_server_outer_keyboard_interrupt(runtime): + with ( + patch.object(runtime, "_validate_auth", side_effect=KeyboardInterrupt()), + patch.object(runtime, "_cleanup", new=AsyncMock()), + ): + result = await runtime._run_server() + assert result.status == UiPathRuntimeStatus.SUCCESSFUL + + +@pytest.mark.asyncio +async def test_run_server_wraps_generic_error(runtime): + with ( + patch.object(runtime, "_validate_auth", side_effect=ValueError("bad")), + patch.object(runtime, "_cleanup", new=AsyncMock()), + ): + with pytest.raises(UiPathMcpRuntimeError): + await runtime._run_server() + + +# --------------------------------------------------------------------------- # +# _register +# --------------------------------------------------------------------------- # +def _tool(name="t", description="d", input_schema=None): + # NB: `name` is a reserved MagicMock ctor kwarg, so it must be set explicitly + # afterwards rather than passed to the constructor. + tool = MagicMock() + tool.name = name + tool.description = description + tool.inputSchema = input_schema + return tool + + +@pytest.mark.asyncio +async def test_register_stdio_success(runtime): + runtime._server.is_streamable_http = False + runtime._job_id = None + runtime._process_key = None # SelfHosted + tools_result = MagicMock(tools=[_tool(input_schema={"type": "object"})]) + session = MagicMock() + session.initialize = AsyncMock() + session.list_tools = AsyncMock(return_value=tools_result) + runtime._uipath.api_client.request_async = AsyncMock() + with ( + patch( + f"{RMODULE}.stdio_client", + return_value=_FakeCM((MagicMock(), MagicMock())), + ), + patch(f"{RMODULE}.ClientSession", return_value=_FakeCM(session)), + ): + await runtime._register() + runtime._uipath.api_client.request_async.assert_awaited_once() + # the registration payload carries the discovered tool with its schema + _args, kwargs = runtime._uipath.api_client.request_async.call_args + payload = kwargs["json"] + assert payload["server"]["Slug"] == runtime.slug + assert payload["tools"][0]["Name"] == "t" + # assert the schema by value, not by exact serialized string + assert json.loads(payload["tools"][0]["InputSchema"]) == {"type": "object"} + + +@pytest.mark.asyncio +async def test_register_stdio_init_failure(runtime): + runtime._server.is_streamable_http = False + runtime._job_id = None + runtime._process_key = None + session = MagicMock() + session.initialize = AsyncMock(side_effect=RuntimeError("init boom")) + with ( + patch( + f"{RMODULE}.stdio_client", + return_value=_FakeCM((MagicMock(), MagicMock())), + ), + patch(f"{RMODULE}.ClientSession", return_value=_FakeCM(session)), + patch.object(runtime, "_on_runtime_abort", new=AsyncMock()), + ): + with pytest.raises(UiPathMcpRuntimeError): + await runtime._register() + + +@pytest.mark.asyncio +async def test_register_http_success_coded(runtime): + runtime._server.is_streamable_http = True + runtime._server.url = "http://localhost:9" + runtime._process_key = CODED_KEY # Coded -> include env + tools_result = MagicMock(tools=[_tool(description=None, input_schema=None)]) + session = MagicMock() + session.initialize = AsyncMock() + session.list_tools = AsyncMock(return_value=tools_result) + runtime._uipath.api_client.request_async = AsyncMock() + with ( + patch.object(runtime, "_start_http_server_process", new=AsyncMock()), + patch.object(runtime, "_wait_for_http_server_ready", new=AsyncMock()), + patch( + "mcp.client.streamable_http.streamable_http_client", + return_value=_FakeCM((MagicMock(), MagicMock(), MagicMock())), + ), + patch(f"{RMODULE}.ClientSession", return_value=_FakeCM(session)), + ): + await runtime._register() + runtime._uipath.api_client.request_async.assert_awaited_once() + # tool with no input schema serializes to an empty-object schema + _args, kwargs = runtime._uipath.api_client.request_async.call_args + payload = kwargs["json"] + assert payload["tools"][0]["Name"] == "t" + assert json.loads(payload["tools"][0]["InputSchema"]) == {} + + +@pytest.mark.asyncio +async def test_register_registration_http_error(runtime): + runtime._server.is_streamable_http = False + runtime._job_id = None + runtime._process_key = None + tools_result = MagicMock(tools=[]) + session = MagicMock() + session.initialize = AsyncMock() + session.list_tools = AsyncMock(return_value=tools_result) + err = httpx.HTTPStatusError( + "e", request=MagicMock(), response=MagicMock(status_code=500, text="server err") + ) + runtime._uipath.api_client.request_async = AsyncMock(side_effect=err) + with ( + patch( + f"{RMODULE}.stdio_client", + return_value=_FakeCM((MagicMock(), MagicMock())), + ), + patch(f"{RMODULE}.ClientSession", return_value=_FakeCM(session)), + ): + with pytest.raises(UiPathMcpRuntimeError): + await runtime._register() + + +# --------------------------------------------------------------------------- # +# _keep_alive +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_keep_alive_pings_and_cancels_sandbox(runtime): + runtime._job_id = "job" # sandboxed + + async def fake_send(method, arguments, on_invocation): + await on_invocation(MagicMock(error=None, result=[])) + + runtime._signalr_client = MagicMock(send=fake_send) + with patch(f"{RMODULE}.asyncio.wait_for", new=AsyncMock(return_value=None)): + await runtime._keep_alive() + assert runtime._cancel_event.is_set() + + +@pytest.mark.asyncio +async def test_keep_alive_error_response_does_not_cancel(runtime): + # A sandboxed runtime with an *empty* session list would normally self-cancel; + # an error response must short-circuit before that check, leaving it running. + runtime._job_id = "job" # sandboxed + + async def fake_send(method, arguments, on_invocation): + await on_invocation(MagicMock(error="ka-boom", result=[])) + + runtime._signalr_client = MagicMock(send=fake_send) + with patch(f"{RMODULE}.asyncio.wait_for", new=AsyncMock(return_value=None)): + await runtime._keep_alive() + assert not runtime._cancel_event.is_set() + + +@pytest.mark.asyncio +async def test_keep_alive_no_signalr_client(runtime, caplog): + # This branch has no observable effect other than the log line (there is no + # client to act on), so a log assertion is the honest check here. + runtime._signalr_client = None + with ( + patch(f"{RMODULE}.asyncio.wait_for", new=AsyncMock(return_value=None)), + caplog.at_level(logging.ERROR), + ): + await runtime._keep_alive() + assert "SignalR client not initialized" in caplog.text + + +@pytest.mark.asyncio +async def test_keep_alive_send_error_is_swallowed(runtime): + send = AsyncMock(side_effect=RuntimeError("send-nope")) + runtime._signalr_client = MagicMock(send=send) + with patch(f"{RMODULE}.asyncio.wait_for", new=AsyncMock(return_value=None)): + # a send failure must not break the heartbeat loop or cancel the runtime + await runtime._keep_alive() + send.assert_awaited_once() + assert not runtime._cancel_event.is_set() + + +@pytest.mark.asyncio +async def test_keep_alive_cancelled(runtime): + runtime._signalr_client = MagicMock(send=AsyncMock()) + with patch( + f"{RMODULE}.asyncio.wait_for", + new=AsyncMock(side_effect=asyncio.CancelledError()), + ): + with pytest.raises(asyncio.CancelledError): + await runtime._keep_alive() diff --git a/tests/cli/test_session_more.py b/tests/cli/test_session_more.py new file mode 100644 index 0000000..8d1abae --- /dev/null +++ b/tests/cli/test_session_more.py @@ -0,0 +1,245 @@ +"""Additional session-server tests: relay loop, consume paths, start/run flows. + +Structural unit tests with the stdio/streamable-http transports mocked; they +assert observable effects (stream wiring, queue draining, error swallowing) +rather than log text. The agenthub wire protocol is exercised by the +integration suite (integration_tests.yml), not here. +""" + +import asyncio +import contextlib +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from mcp.types import ( + JSONRPCMessage, + JSONRPCRequest, + JSONRPCResponse, +) + +from uipath_mcp._cli._runtime._session import ( + StdioSessionServer, + StreamableHttpSessionServer, +) + +SMODULE = "uipath_mcp._cli._runtime._session" + + +def _make_server_config(url: str | None = None) -> MagicMock: + cfg = MagicMock() + cfg.command = "cmd" + cfg.args = [] + cfg.env = {} + cfg.url = url + return cfg + + +@pytest.fixture +def uipath_mock() -> MagicMock: + m = MagicMock() + m.api_client.request_async = AsyncMock() + return m + + +@pytest.fixture +def stdio_session(uipath_mock) -> StdioSessionServer: + return StdioSessionServer(_make_server_config(), "slug", "sess-id", uipath_mock) + + +@pytest.fixture +def http_session(uipath_mock) -> StreamableHttpSessionServer: + return StreamableHttpSessionServer( + _make_server_config(url="https://x"), "slug", "sess-id", uipath_mock + ) + + +def test_is_response_without_root_returns_false(stdio_session): + assert stdio_session._is_response(object()) is False + + +@pytest.mark.asyncio +async def test_stop_logs_unexpected_error(stdio_session): + """stop() swallows and logs a non-cancel/timeout error from task teardown.""" + task = MagicMock() + task.done.return_value = False + stdio_session._run_task = task + # Force wait_for to raise a non-cancel/timeout error so the generic `except` + # branch runs; patching the primitive is the least-bad way to reach it. + with ( + patch(f"{SMODULE}.asyncio.shield", return_value=MagicMock()), + patch( + f"{SMODULE}.asyncio.wait_for", new=AsyncMock(side_effect=RuntimeError("x")) + ), + ): + await stdio_session.stop() + assert stdio_session._run_task is None + + +@pytest.mark.asyncio +async def test_relay_messages_routes_and_reports_errors(stdio_session): + s = stdio_session + s._send_message = AsyncMock() + s._last_request_id = "reqLast" + s._last_message_id = "9" + s._active_requests = {"5": "reqMapped"} + + resp_mapped = MagicMock() + resp_mapped.message = JSONRPCMessage( + JSONRPCResponse(jsonrpc="2.0", id=5, result={}) + ) + resp_unmapped = MagicMock() + resp_unmapped.message = JSONRPCMessage( + JSONRPCResponse(jsonrpc="2.0", id=99, result={}) + ) + req_msg = MagicMock() + req_msg.message = JSONRPCMessage(JSONRPCRequest(jsonrpc="2.0", id=1, method="m")) + error_item = ValueError("received-as-error") + + queue = [resp_mapped, resp_unmapped, req_msg, error_item] + + async def fake_receive(): + if queue: + return queue.pop(0) + s._read_stream = None + raise RuntimeError("stop loop") + + s._read_stream = MagicMock() + s._read_stream.receive = fake_receive + + await s._relay_messages() + + # mapped response consumed the stored mapping + assert "5" not in s._active_requests + # mapped + unmapped + request + final error report + assert s._send_message.await_count >= 4 + + +@pytest.mark.asyncio +async def test_consume_messages_send_error_does_not_hang(stdio_session): + s = stdio_session + s._write_stream = MagicMock() + s._write_stream.send = AsyncMock(side_effect=RuntimeError("send fail")) + await s._message_queue.put( + JSONRPCMessage(JSONRPCRequest(jsonrpc="2.0", id=1, method="m")) + ) + task = asyncio.create_task(s._consume_messages()) + await asyncio.sleep(0.05) + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + # a failing send still marks the message done -> the consumer never blocks + assert s._message_queue.empty() + + +@pytest.mark.asyncio +async def test_consume_messages_drains_queue_on_cancel(stdio_session): + s = stdio_session + block = asyncio.Event() + + async def slow_send(_message): + await block.wait() + + s._write_stream = MagicMock() + s._write_stream.send = slow_send + for i in range(2): + await s._message_queue.put( + JSONRPCMessage(JSONRPCRequest(jsonrpc="2.0", id=i, method="m")) + ) + task = asyncio.create_task(s._consume_messages()) + await asyncio.sleep(0.05) # first message in-flight, second still queued + task.cancel() + await task # CancelledError is handled internally; drains remaining + assert s._message_queue.empty() + + +@pytest.mark.asyncio +async def test_stdio_start_error_calls_stop_and_raises(stdio_session): + with patch(f"{SMODULE}.StdioServerParameters", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError, match="boom"): + await stdio_session.start() + + +@pytest.mark.asyncio +async def test_http_start_error_calls_stop_and_raises(http_session): + with patch(f"{SMODULE}.asyncio.create_task", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError, match="boom"): + await http_session.start() + + +class _FakeCM: + def __init__(self, value): + self._value = value + + async def __aenter__(self): + return self._value + + async def __aexit__(self, *_a): + return False + + +@pytest.mark.asyncio +async def test_stdio_run_server_success(stdio_session): + s = stdio_session + read, write = MagicMock(), MagicMock() + with ( + patch(f"{SMODULE}.stdio_client", return_value=_FakeCM((read, write))), + patch.object(s, "_relay_messages", new=AsyncMock()) as relay, + ): + await s._run_server(MagicMock()) + # streams from stdio_client were wired in and the relay loop was driven + relay.assert_awaited_once() + assert s._read_stream is read + assert s._write_stream is write + assert s._server_stderr_output is not None + + +@pytest.mark.asyncio +async def test_stdio_run_server_handles_exception_group(stdio_session): + s = stdio_session + with ( + patch( + f"{SMODULE}.stdio_client", + return_value=_FakeCM((MagicMock(), MagicMock())), + ), + patch.object( + s, "_relay_messages", new=AsyncMock(side_effect=RuntimeError("relay boom")) + ), + ): + # the except* group handler swallows the error; finally still captures stderr + await s._run_server(MagicMock()) + assert s._server_stderr_output is not None + + +@pytest.mark.asyncio +async def test_http_run_session_success(http_session): + s = http_session + read, write = MagicMock(), MagicMock() + with ( + patch( + f"{SMODULE}.streamable_http_client", + return_value=_FakeCM((read, write, MagicMock())), + ), + patch.object(s, "_relay_messages", new=AsyncMock()) as relay, + ): + await s._run_http_session() + relay.assert_awaited_once() + assert s._read_stream is read + assert s._write_stream is write + + +@pytest.mark.asyncio +async def test_http_run_session_swallows_exception(http_session): + s = http_session + read = MagicMock() + with ( + patch( + f"{SMODULE}.streamable_http_client", + return_value=_FakeCM((read, MagicMock(), MagicMock())), + ), + patch.object( + s, "_relay_messages", new=AsyncMock(side_effect=RuntimeError("http boom")) + ), + ): + await s._run_http_session() # exception is caught, not propagated + # the connection was established (streams wired) before the failure + assert s._read_stream is read diff --git a/tests/cli/test_token_refresh_more.py b/tests/cli/test_token_refresh_more.py new file mode 100644 index 0000000..669af2c --- /dev/null +++ b/tests/cli/test_token_refresh_more.py @@ -0,0 +1,238 @@ +"""Additional TokenRefresher tests covering loop, refresh flows, and lifecycle. + +Identity/OAuth endpoints and the HTTP client are mocked; these assert the +refresher's own decision logic (strategy selection, retry/cancel, propagation), +not the real identity service contract. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from uipath.platform.common import TokenData +from uipath.platform.common._config import UiPathApiConfig + +from uipath_mcp._cli._runtime._token_refresh import ( + AuthStrategy, + TokenRefresher, +) + +MODULE = "uipath_mcp._cli._runtime._token_refresh" + + +def _make_uipath_mock( + base_url: str = "https://cloud.uipath.com/org/tenant", + secret: str = "mock_token", +) -> MagicMock: + mock = MagicMock() + mock._config = UiPathApiConfig(base_url=base_url, secret=secret) + return mock + + +def _client_creds_refresher( + monkeypatch, *, token_url="https://id/token" +) -> TokenRefresher: + monkeypatch.setenv("UIPATH_CLIENT_ID", "cid") + monkeypatch.setenv("UIPATH_CLIENT_SECRET", "csecret") + with patch(f"{MODULE}.build_service_url", return_value=token_url): + return TokenRefresher(_make_uipath_mock()) + + +def test_client_credentials_without_token_url_disables_refresh(monkeypatch): + """CLIENT_CREDENTIALS with unresolvable token URL falls back to NONE.""" + monkeypatch.setenv("UIPATH_CLIENT_ID", "cid") + monkeypatch.setenv("UIPATH_CLIENT_SECRET", "csecret") + with patch(f"{MODULE}.build_service_url", side_effect=RuntimeError("no url")): + refresher = TokenRefresher(_make_uipath_mock()) + assert refresher.strategy == AuthStrategy.NONE + assert refresher._token_url is None + + +@pytest.mark.asyncio +async def test_start_and_stop_lifecycle(monkeypatch): + """start() spawns a task for a live strategy and stop() cancels it.""" + refresher = _client_creds_refresher(monkeypatch) + assert refresher.strategy == AuthStrategy.CLIENT_CREDENTIALS + + with patch.object(refresher, "_refresh_loop", new=AsyncMock()): + refresher.start() + assert refresher._refresh_task is not None + await refresher.stop() + + assert refresher._refresh_task is None + + +@pytest.mark.asyncio +async def test_wait_for_cancel_times_out(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + # cancel event never set -> times out and returns False + assert await refresher._wait_for_cancel(0.01) is False + + +@pytest.mark.asyncio +async def test_wait_for_cancel_returns_true_when_set(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + refresher._cancel_event.set() + assert await refresher._wait_for_cancel(1.0) is True + + +@pytest.mark.asyncio +async def test_refresh_loop_success_then_cancel(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + + async def fake_try() -> bool: + refresher._cancel_event.set() + return True + + with ( + patch.object(refresher, "_seconds_until_refresh", return_value=0), + patch.object(refresher, "_try_refresh", side_effect=fake_try) as tr, + ): + await refresher._refresh_loop() + + tr.assert_awaited_once() + assert refresher._cancel_event.is_set() + + +@pytest.mark.asyncio +async def test_refresh_loop_waits_then_cancel(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + with ( + patch.object(refresher, "_seconds_until_refresh", return_value=5), + patch.object( + refresher, "_wait_for_cancel", new=AsyncMock(return_value=True) + ) as wc, + patch.object(refresher, "_try_refresh", new=AsyncMock()) as tr, + ): + await refresher._refresh_loop() + + wc.assert_awaited_once_with(5) + tr.assert_not_awaited() # broke during the pre-refresh wait, never refreshed + + +@pytest.mark.asyncio +async def test_refresh_loop_all_attempts_fail_then_break(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + with ( + patch.object(refresher, "_seconds_until_refresh", return_value=0), + patch.object( + refresher, "_try_refresh", new=AsyncMock(return_value=False) + ) as tr, + patch.object( + refresher, "_wait_for_cancel", new=AsyncMock(return_value=True) + ) as wc, + ): + await refresher._refresh_loop() + + tr.assert_awaited_once() + # after all attempts fail it waits the fallback interval, then breaks + wc.assert_awaited_once_with(60) + + +@pytest.mark.asyncio +async def test_refresh_loop_propagates_cancelled(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + with ( + patch.object(refresher, "_seconds_until_refresh", return_value=5), + patch.object( + refresher, + "_wait_for_cancel", + new=AsyncMock(side_effect=asyncio.CancelledError()), + ), + pytest.raises(asyncio.CancelledError), + ): + await refresher._refresh_loop() + + +@pytest.mark.asyncio +async def test_try_refresh_oauth_success(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + refresher._strategy = AuthStrategy.OAUTH + token = TokenData(access_token="new") + with ( + patch.object(refresher, "_refresh_oauth", new=AsyncMock(return_value=token)), + patch.object(refresher, "_propagate_token") as prop, + ): + assert await refresher._try_refresh() is True + prop.assert_called_once_with(token) + + +@pytest.mark.asyncio +async def test_try_refresh_retries_then_cancel(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + with ( + patch.object( + refresher, + "_refresh_client_credentials", + new=AsyncMock( + side_effect=httpx.HTTPStatusError( + "x", request=MagicMock(), response=MagicMock(status_code=401) + ) + ), + ), + patch.object(refresher, "_wait_for_cancel", new=AsyncMock(return_value=True)), + ): + assert await refresher._try_refresh() is False + + +@pytest.mark.asyncio +async def test_refresh_oauth_flow(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + auth = MagicMock() + auth.refresh_token = "rt" + token = TokenData(access_token="oauth-token") + identity = MagicMock() + identity.refresh_access_token_async = AsyncMock(return_value=token) + with ( + patch(f"{MODULE}.get_auth_data", return_value=auth), + patch( + f"{MODULE}.OidcUtils.get_auth_config", + new=AsyncMock(return_value={"client_id": "abc"}), + ), + patch(f"{MODULE}.IdentityService", return_value=identity), + patch(f"{MODULE}.update_auth_file", side_effect=RuntimeError("write fail")), + ): + result = await refresher._refresh_oauth() + assert result is token + + +@pytest.mark.asyncio +async def test_refresh_oauth_without_refresh_token_raises(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + auth = MagicMock() + auth.refresh_token = None + with patch(f"{MODULE}.get_auth_data", return_value=auth): + with pytest.raises(ValueError, match="refresh_token"): + await refresher._refresh_oauth() + + +@pytest.mark.asyncio +async def test_refresh_client_credentials_flow(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + + response = MagicMock() + response.raise_for_status = MagicMock() + response.json = MagicMock(return_value={"access_token": "cc-token"}) + client = MagicMock() + client.post = AsyncMock(return_value=response) + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=client) + cm.__aexit__ = AsyncMock(return_value=False) + + with ( + patch(f"{MODULE}.httpx.AsyncClient", return_value=cm), + patch(f"{MODULE}.get_httpx_client_kwargs", return_value={}), + ): + result = await refresher._refresh_client_credentials() + + assert result.access_token == "cc-token" + client.post.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_refresh_client_credentials_without_url_raises(monkeypatch): + refresher = _client_creds_refresher(monkeypatch) + refresher._token_url = None + with pytest.raises(RuntimeError, match="token_url"): + await refresher._refresh_client_credentials() diff --git a/tests/test_diagnose.py b/tests/test_diagnose.py new file mode 100644 index 0000000..48afbbc --- /dev/null +++ b/tests/test_diagnose.py @@ -0,0 +1,147 @@ +"""Tests for the binary diagnostics helper.""" + +import subprocess +from unittest.mock import patch + +from uipath_mcp._cli._utils._diagnose import diagnose_binary + +MODULE = "uipath_mcp._cli._utils._diagnose" + + +def test_missing_file_returns_error(): + with patch(f"{MODULE}.os.path.exists", return_value=False): + result = diagnose_binary("/no/such/file") + assert result == "Error: /no/such/file does not exist" + + +def test_non_linux_file_command_success(): + """Windows/mac path: file type resolved, Linux ELF block skipped.""" + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="AMD64"), + patch(f"{MODULE}.platform.system", return_value="Windows"), + patch( + f"{MODULE}.subprocess.check_output", + return_value="PE32+ executable (console) x86-64\n", + ), + ): + result = diagnose_binary("/bin/app.exe") + assert "incompatible" in result + + +def test_file_command_not_found(): + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="x86_64"), + patch(f"{MODULE}.platform.system", return_value="Windows"), + patch( + f"{MODULE}.subprocess.check_output", + side_effect=FileNotFoundError(), + ), + ): + result = diagnose_binary("/bin/app") + assert "incompatible" in result + + +def test_file_command_process_error(): + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="x86_64"), + patch(f"{MODULE}.platform.system", return_value="Darwin"), + patch( + f"{MODULE}.subprocess.check_output", + side_effect=subprocess.CalledProcessError(1, "file"), + ), + ): + result = diagnose_binary("/bin/app") + assert "incompatible" in result + + +def test_linux_elf_x86_64_arch_mismatch(): + """ELF x86-64 binary on a non-x86_64 host reports the arch mismatch.""" + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="aarch64"), + patch(f"{MODULE}.platform.system", return_value="Linux"), + patch( + f"{MODULE}.subprocess.check_output", + side_effect=[ + "ELF 64-bit LSB executable, x86-64\n", + " Machine: Advanced Micro Devices X86-64\n", + ], + ), + ): + result = diagnose_binary("/bin/app") + assert "x86_64 architecture" in result + assert "aarch64" in result + + +def test_linux_readelf_no_machine_line(): + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="x86_64"), + patch(f"{MODULE}.platform.system", return_value="Linux"), + patch( + f"{MODULE}.subprocess.check_output", + side_effect=[ + "ELF 64-bit LSB executable, x86-64\n", + "no arch details here\n", + ], + ), + ): + result = diagnose_binary("/bin/app") + # x86_64 host + x86-64 binary => no mismatch, falls through to generic msg + assert "incompatible" in result + + +def test_linux_readelf_process_error(): + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="x86_64"), + patch(f"{MODULE}.platform.system", return_value="Linux"), + patch( + f"{MODULE}.subprocess.check_output", + side_effect=[ + "ELF 64-bit LSB executable, x86-64\n", + subprocess.CalledProcessError(1, "readelf"), + ], + ), + ): + result = diagnose_binary("/bin/app") + assert "incompatible" in result + + +def test_linux_readelf_not_found(): + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="x86_64"), + patch(f"{MODULE}.platform.system", return_value="Linux"), + patch( + f"{MODULE}.subprocess.check_output", + side_effect=[ + "ELF 64-bit LSB executable, x86-64\n", + FileNotFoundError(), + ], + ), + ): + result = diagnose_binary("/bin/app") + assert "incompatible" in result + + +def test_linux_elf_arm_mismatch(): + """ARM binary on an x86_64 host reports the ARM mismatch.""" + with ( + patch(f"{MODULE}.os.path.exists", return_value=True), + patch(f"{MODULE}.platform.machine", return_value="x86_64"), + patch(f"{MODULE}.platform.system", return_value="Linux"), + patch( + f"{MODULE}.subprocess.check_output", + side_effect=[ + "ELF 32-bit LSB executable, ARM\n", + " Machine: ARM\n", + ], + ), + ): + result = diagnose_binary("/bin/app") + assert "ARM architecture" in result + assert "x86_64" in result