[WIP] feat(MockAlwaysAwaitable): mockables are always awaitable

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]))
 

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__,

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,