diff --git a/samples/calculator/main.py b/samples/calculator/main.py index 31eb43af0..adf4a1c53 100644 --- a/samples/calculator/main.py +++ b/samples/calculator/main.py @@ -35,7 +35,7 @@ class Wrapper: @traced() @mockable(example_calls=GET_RANDOM_OPERATOR_EXAMPLES) -async def get_random_operator() -> Wrapper: +def get_random_operator() -> Wrapper: """Get a random operator.""" return Wrapper(result=random.choice([Operator.ADD, Operator.SUBTRACT, Operator.MULTIPLY, Operator.DIVIDE])) diff --git a/src/uipath/eval/mocks/mockable.py b/src/uipath/eval/mocks/mockable.py index a8e03e625..adc9bf4bd 100644 --- a/src/uipath/eval/mocks/mockable.py +++ b/src/uipath/eval/mocks/mockable.py @@ -5,7 +5,7 @@ import inspect import logging import threading -from typing import Any, List, Optional +from typing import Any, Awaitable, Callable, List, Optional, ParamSpec, TypeVar from pydantic import TypeAdapter from pydantic_function_models import ValidatedFunction # type: ignore[import-untyped] @@ -28,7 +28,13 @@ def run_coroutine(coro): return future.result() -def mocked_response_decorator(func, params: dict[str, Any]): +T = ParamSpec("T") +R = TypeVar("R") + + +def mocked_response_decorator( + func: Callable[T, R | Awaitable[R]], params: dict[str, Any] +) -> Callable[T, Awaitable[R]]: """Mocked response decorator.""" async def mock_response_generator(*args, **kwargs): @@ -39,23 +45,16 @@ async def mock_response_generator(*args, **kwargs): mocked_response = TypeAdapter(return_type).validate_python(mocked_response) return mocked_response - is_async = inspect.iscoroutinefunction(func) - if is_async: - - @functools.wraps(func) - async def decorated_func(*args, **kwargs): - try: - return await mock_response_generator(*args, **kwargs) - except UiPathNoMockFoundError: - return await func(*args, **kwargs) - else: - - @functools.wraps(func) - def decorated_func(*args, **kwargs): - try: - return run_coroutine(mock_response_generator(*args, **kwargs)) - except UiPathNoMockFoundError: - return func(*args, **kwargs) + @functools.wraps(func) + async def decorated_func(*args, **kwargs) -> R: + try: + return await mock_response_generator(*args, **kwargs) + except UiPathNoMockFoundError: + result = func(*args, **kwargs) + if isinstance(result, Awaitable): + return await result + else: + return result return decorated_func @@ -86,10 +85,10 @@ def mockable( output_schema: Optional[dict[str, Any]] = None, example_calls: Optional[List[ExampleCall]] = None, **kwargs, -): +) -> Callable[[Callable[T, R | Awaitable[R]]], Callable[T, Awaitable[R]]]: """Decorate a function to be a mockable.""" - def decorator(func): + def decorator(func: Callable[T, R | Awaitable[R]]) -> Callable[T, Awaitable[R]]: params = { "name": name or func.__name__, "description": description or func.__doc__, diff --git a/tests/cli/eval/mocks/test_mocks.py b/tests/cli/eval/mocks/test_mocks.py index d374bb74c..66ead9473 100644 --- a/tests/cli/eval/mocks/test_mocks.py +++ b/tests/cli/eval/mocks/test_mocks.py @@ -1,4 +1,5 @@ -from typing import Any +import asyncio +from typing import Any, Awaitable import pytest from _pytest.monkeypatch import MonkeyPatch @@ -14,63 +15,17 @@ from uipath.eval.mocks import mockable -def test_mockito_mockable_sync(): - # Arrange +def test_mockable_is_always_awaitable(): @mockable() - def foo(*args, **kwargs): - raise NotImplementedError() + def foo(): + return "success" - @mockable() - def foofoo(*args, **kwargs): - raise NotImplementedError() - - evaluation_item: dict[str, Any] = { - "id": "evaluation-id", - "name": "Mock foo", - "inputs": {}, - "expectedOutput": {}, - "expectedAgentBehavior": "", - "mockingStrategy": { - "type": "mockito", - "behaviors": [ - { - "function": "foo", - "arguments": {"args": [], "kwargs": {}}, - "then": [ - {"type": "return", "value": "bar1"}, - {"type": "return", "value": "bar2"}, - ], - } - ], - }, - "evalSetId": "eval-set-id", - "createdAt": "2025-09-04T18:54:58.378Z", - "updatedAt": "2025-09-04T18:55:55.416Z", - } - evaluation = EvaluationItem(**evaluation_item) - assert isinstance(evaluation.mocking_strategy, MockitoMockingStrategy) - - # Act & Assert - set_evaluation_item(evaluation) - assert foo() == "bar1" - assert foo() == "bar2" - assert foo() == "bar2" - - with pytest.raises(UiPathMockResponseGenerationError): - assert foo(x=1) - - with pytest.raises(NotImplementedError): - assert foofoo() - - evaluation.mocking_strategy.behaviors[0].arguments.kwargs = {"x": 1} - set_evaluation_item(evaluation) - assert foo(x=1) == "bar1" - - evaluation.mocking_strategy.behaviors[0].arguments.kwargs = { - "x": {"_target_": "mockito.any"} - } - set_evaluation_item(evaluation) - assert foo(x=2) == "bar1" + awaitable_result = foo() + assert isinstance(awaitable_result, Awaitable) + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + result = loop.run_until_complete(awaitable_result) + assert result == "success" @pytest.mark.asyncio @@ -133,17 +88,18 @@ async def foofoo(*args, **kwargs): assert await foo(x=2) == "bar1" -def test_llm_mockable_sync(httpx_mock: HTTPXMock, monkeypatch: MonkeyPatch): +@pytest.mark.asyncio +async def test_llm_mockable_async(httpx_mock: HTTPXMock, monkeypatch: MonkeyPatch): monkeypatch.setenv("UIPATH_URL", "https://example.com") monkeypatch.setenv("UIPATH_ACCESS_TOKEN", "1234567890") # Arrange @mockable() - def foo(*args, **kwargs): + async def foo(*args, **kwargs): raise NotImplementedError() @mockable() - def foofoo(*args, **kwargs): + async def foofoo(*args, **kwargs): raise NotImplementedError() evaluation_item: dict[str, Any] = { @@ -163,6 +119,7 @@ def foofoo(*args, **kwargs): } evaluation = EvaluationItem(**evaluation_item) assert isinstance(evaluation.mocking_strategy, LLMMockingStrategy) + httpx_mock.add_response( url="https://example.com/agenthub_/llm/api/capabilities", status_code=200, @@ -173,80 +130,6 @@ def foofoo(*args, **kwargs): status_code=200, json={}, ) - - httpx_mock.add_response( - url="https://example.com/api/chat/completions?api-version=2024-08-01-preview", - status_code=200, - json={ - "id": "response-id", - "object": "", - "created": 0, - "model": "model", - "choices": [ - { - "index": 0, - "message": { - "role": "ai", - "content": '{"response": "bar1"}', - "tool_calls": None, - }, - "finish_reason": "EOS", - } - ], - "usage": { - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - }, - ) - # Act & Assert - set_evaluation_item(evaluation) - - assert foo() == "bar1" - with pytest.raises(NotImplementedError): - assert foofoo() - httpx_mock.add_response( - url="https://example.com/api/chat/completions?api-version=2024-08-01-preview", - status_code=200, - json={}, - ) - with pytest.raises(UiPathMockResponseGenerationError): - assert foo() - - -@pytest.mark.asyncio -async def test_llm_mockable_async(httpx_mock: HTTPXMock, monkeypatch: MonkeyPatch): - monkeypatch.setenv("UIPATH_URL", "https://example.com") - monkeypatch.setenv("UIPATH_ACCESS_TOKEN", "1234567890") - - # Arrange - @mockable() - async def foo(*args, **kwargs): - raise NotImplementedError() - - @mockable() - async def foofoo(*args, **kwargs): - raise NotImplementedError() - - evaluation_item: dict[str, Any] = { - "id": "evaluation-id", - "name": "Mock foo", - "inputs": {}, - "expectedOutput": {}, - "expectedAgentBehavior": "", - "mockingStrategy": { - "type": "llm", - "prompt": "response is 'bar1'", - "toolsToSimulate": [{"name": "foo"}], - }, - "evalSetId": "eval-set-id", - "createdAt": "2025-09-04T18:54:58.378Z", - "updatedAt": "2025-09-04T18:55:55.416Z", - } - evaluation = EvaluationItem(**evaluation_item) - assert isinstance(evaluation.mocking_strategy, LLMMockingStrategy) - httpx_mock.add_response( url="https://example.com/api/chat/completions?api-version=2024-08-01-preview", status_code=200,