hatchet-sdk 1.15.3__py3-none-any.whl → 1.16.1__py3-none-any.whl

hatchet_sdk/features/cel.py

@@ -0,0 +1,99 @@
+import asyncio
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+from hatchet_sdk.clients.rest.api.cel_api import CELApi
+from hatchet_sdk.clients.rest.api_client import ApiClient
+from hatchet_sdk.clients.rest.models.v1_cel_debug_request import V1CELDebugRequest
+from hatchet_sdk.clients.rest.models.v1_cel_debug_response_status import (
+    V1CELDebugResponseStatus,
+)
+from hatchet_sdk.clients.v1.api_client import BaseRestClient, retry
+from hatchet_sdk.utils.typing import JSONSerializableMapping
+
+
+class CELSuccess(BaseModel):
+    status: Literal["success"] = "success"
+    output: bool
+
+
+class CELFailure(BaseModel):
+    status: Literal["failure"] = "failure"
+    error: str
+
+
+class CELEvaluationResult(BaseModel):
+    result: CELSuccess | CELFailure = Field(discriminator="status")
+
+
+class CELClient(BaseRestClient):
+    """
+    The CEL client is a client for debugging CEL expressions within Hatchet
+    """
+
+    def _ca(self, client: ApiClient) -> CELApi:
+        return CELApi(client)
+
+    @retry
+    def debug(
+        self,
+        expression: str,
+        input: JSONSerializableMapping,
+        additional_metadata: JSONSerializableMapping | None = None,
+        filter_payload: JSONSerializableMapping | None = None,
+    ) -> CELEvaluationResult:
+        """
+        Debug a CEL expression with the provided input, filter payload, and optional metadata. Useful for testing and validating CEL expressions and debugging issues in production.
+
+        :param expression: The CEL expression to debug.
+        :param input: The input, which simulates the workflow run input.
+        :param additional_metadata: Additional metadata, which simulates metadata that could be sent with an event or a workflow run
+        :param filter_payload: The filter payload, which simulates a payload set on a previous-created filter
+
+        :raises ValueError: If no response is received from the CEL debug API.
+
+        :return: A V1CELDebugErrorResponse or V1CELDebugSuccessResponse containing the result of the debug operation.
+        """
+        request = V1CELDebugRequest(
+            expression=expression,
+            input=input,
+            additionalMetadata=additional_metadata,
+            filterPayload=filter_payload,
+        )
+        with self.client() as client:
+            result = self._ca(client).v1_cel_debug(
+                tenant=self.client_config.tenant_id, v1_cel_debug_request=request
+            )
+
+            if result.status == V1CELDebugResponseStatus.ERROR:
+                if result.error is None:
+                    raise ValueError("No error message received from CEL debug API.")
+
+                return CELEvaluationResult(result=CELFailure(error=result.error))
+
+            if result.output is None:
+                raise ValueError("No output received from CEL debug API.")
+
+            return CELEvaluationResult(result=CELSuccess(output=result.output))
+
+    async def aio_debug(
+        self,
+        expression: str,
+        input: JSONSerializableMapping,
+        additional_metadata: JSONSerializableMapping | None = None,
+        filter_payload: JSONSerializableMapping | None = None,
+    ) -> CELEvaluationResult:
+        """
+        Debug a CEL expression with the provided input, filter payload, and optional metadata. Useful for testing and validating CEL expressions and debugging issues in production.
+
+        :param expression: The CEL expression to debug.
+        :param input: The input, which simulates the workflow run input.
+        :param additional_metadata: Additional metadata, which simulates metadata that could be sent with an event or a workflow run
+        :param filter_payload: The filter payload, which simulates a payload set on a previous-created filter
+
+        :return: A V1CELDebugErrorResponse or V1CELDebugSuccessResponse containing the result of the debug operation.
+        """
+        return await asyncio.to_thread(
+            self.debug, expression, input, additional_metadata, filter_payload
+        )

hatchet_sdk/features/runs.py

@@ -119,6 +119,26 @@ class RunsClient(BaseRestClient):
     def _ta(self, client: ApiClient) -> TaskApi:
         return TaskApi(client)
 
+    @retry
+    def get_task_run(self, task_run_id: str) -> V1TaskSummary:
+        """
+        Get task run details for a given task run ID.
+
+        :param task_run_id: The ID of the task run to retrieve details for.
+        :return: Task run details for the specified task run ID.
+        """
+        with self.client() as client:
+            return self._ta(client).v1_task_get(task_run_id)
+
+    async def aio_get_task_run(self, task_run_id: str) -> V1TaskSummary:
+        """
+        Get task run details for a given task run ID.
+
+        :param task_run_id: The ID of the task run to retrieve details for.
+        :return: Task run details for the specified task run ID.
+        """
+        return await asyncio.to_thread(self.get_task_run, task_run_id)
+
     @retry
     def get(self, workflow_run_id: str) -> V1WorkflowRunDetails:
         """
@@ -148,7 +168,7 @@ class RunsClient(BaseRestClient):
         :return: The task status
         """
         with self.client() as client:
-            return self._wra(client).v1_workflow_run_get_status(str(workflow_run_id))
+            return self._wra(client).v1_workflow_run_get_status(workflow_run_id)
 
     async def aio_get_status(self, workflow_run_id: str) -> V1TaskStatus:
         """

hatchet_sdk/hatchet.py

@@ -12,6 +12,7 @@ from hatchet_sdk.clients.events import EventClient
 from hatchet_sdk.clients.listeners.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.rest.models.tenant_version import TenantVersion
 from hatchet_sdk.config import ClientConfig
+from hatchet_sdk.features.cel import CELClient
 from hatchet_sdk.features.cron import CronClient
 from hatchet_sdk.features.filters import FiltersClient
 from hatchet_sdk.features.logs import LogsClient
@@ -66,6 +67,13 @@ class Hatchet:
                 "🚨⚠️‼️ YOU ARE USING A V0 ENGINE WITH A V1 SDK, WHICH IS NOT SUPPORTED. PLEASE UPGRADE YOUR ENGINE TO V1.🚨⚠️‼️"
             )
 
+    @property
+    def cel(self) -> CELClient:
+        """
+        The CEL client is a client for interacting with Hatchet's CEL API.
+        """
+        return self._client.cel
+
     @property
     def cron(self) -> CronClient:
         """

hatchet_sdk/opentelemetry/instrumentor.py

@@ -64,7 +64,7 @@ OTEL_TRACEPARENT_KEY = "traceparent"
 
 def create_traceparent() -> str | None:
     logger.warning(
-        "As of SDK version 1.11.0, you no longer need to call `create_traceparent` manually. The traceparent will be automatically created by the instrumentor and injected into the metadata of actions and events when appropriate. This method will be removed in a future version.",
+        "as of SDK version 1.11.0, you no longer need to call `create_traceparent` manually. The traceparent will be automatically created by the instrumentor and injected into the metadata of actions and events when appropriate. This method will be removed in a future version.",
     )
     return _create_traceparent()
 
@@ -91,7 +91,7 @@ def parse_carrier_from_metadata(
     metadata: JSONSerializableMapping | None,
 ) -> Context | None:
     logger.warning(
-        "As of SDK version 1.11.0, you no longer need to call `parse_carrier_from_metadata` manually. This method will be removed in a future version.",
+        "as of SDK version 1.11.0, you no longer need to call `parse_carrier_from_metadata` manually. This method will be removed in a future version.",
     )
 
     return _parse_carrier_from_metadata(metadata)
@@ -133,7 +133,7 @@ def inject_traceparent_into_metadata(
     metadata: dict[str, str], traceparent: str | None = None
 ) -> dict[str, str]:
     logger.warning(
-        "As of SDK version 1.11.0, you no longer need to call `inject_traceparent_into_metadata` manually. The traceparent will automatically be injected by the instrumentor. This method will be removed in a future version.",
+        "as of SDK version 1.11.0, you no longer need to call `inject_traceparent_into_metadata` manually. The traceparent will automatically be injected by the instrumentor. This method will be removed in a future version.",
     )
 
     return _inject_traceparent_into_metadata(metadata, traceparent)

hatchet_sdk/runnables/contextvars.py

@@ -4,6 +4,7 @@ from collections import Counter
 from contextvars import ContextVar
 
 from hatchet_sdk.runnables.action import ActionKey
+from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 ctx_workflow_run_id: ContextVar[str | None] = ContextVar(
     "ctx_workflow_run_id", default=None
@@ -13,6 +14,9 @@ ctx_action_key: ContextVar[ActionKey | None] = ContextVar(
 )
 ctx_step_run_id: ContextVar[str | None] = ContextVar("ctx_step_run_id", default=None)
 ctx_worker_id: ContextVar[str | None] = ContextVar("ctx_worker_id", default=None)
+ctx_additional_metadata: ContextVar[JSONSerializableMapping | None] = ContextVar(
+    "ctx_additional_metadata", default=None
+)
 
 workflow_spawn_indices = Counter[ActionKey]()
 spawn_index_lock = asyncio.Lock()

hatchet_sdk/runnables/task.py

@@ -11,6 +11,7 @@ from hatchet_sdk.conditions import (
     flatten_conditions,
 )
 from hatchet_sdk.context.context import Context, DurableContext
+from hatchet_sdk.context.worker_context import WorkerContext
 from hatchet_sdk.contracts.v1.shared.condition_pb2 import TaskConditions
 from hatchet_sdk.contracts.v1.workflows_pb2 import (
     CreateTaskOpts,
@@ -19,6 +20,7 @@ from hatchet_sdk.contracts.v1.workflows_pb2 import (
 )
 from hatchet_sdk.runnables.types import (
     ConcurrencyExpression,
+    EmptyModel,
     R,
     StepType,
     TWorkflowInput,
@@ -30,9 +32,11 @@ from hatchet_sdk.utils.timedelta_to_expression import Duration, timedelta_to_exp
 from hatchet_sdk.utils.typing import (
     AwaitableLike,
     CoroutineLike,
+    JSONSerializableMapping,
     TaskIOValidator,
     is_basemodel_subclass,
 )
+from hatchet_sdk.worker.runner.utils.capture_logs import AsyncLogSender
 
 if TYPE_CHECKING:
     from hatchet_sdk.runnables.workflow import Workflow
@@ -186,3 +190,132 @@ class Task(Generic[TWorkflowInput, R]):
             sleep_conditions=sleep_conditions,
             user_event_conditions=user_events,
         )
+
+    def _create_mock_context(
+        self,
+        input: TWorkflowInput | None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan_context: Any = None,
+    ) -> Context | DurableContext:
+        from hatchet_sdk.runnables.action import Action, ActionPayload, ActionType
+
+        additional_metadata = additional_metadata or {}
+        parent_outputs = parent_outputs or {}
+
+        if input is None:
+            input = cast(TWorkflowInput, EmptyModel())
+
+        action_payload = ActionPayload(input=input.model_dump(), parents=parent_outputs)
+
+        action = Action(
+            tenant_id=self.workflow.client.config.tenant_id,
+            worker_id="mock-worker-id",
+            workflow_run_id="mock-workflow-run-id",
+            get_group_key_run_id="mock-get-group-key-run-id",
+            job_id="mock-job-id",
+            job_name="mock-job-name",
+            job_run_id="mock-job-run-id",
+            step_id="mock-step-id",
+            step_run_id="mock-step-run-id",
+            action_id="mock:action",
+            action_payload=action_payload,
+            action_type=ActionType.START_STEP_RUN,
+            retry_count=retry_count,
+            additional_metadata=additional_metadata,
+            child_workflow_index=None,
+            child_workflow_key=None,
+            parent_workflow_run_id=None,
+            priority=1,
+            workflow_version_id="mock-workflow-version-id",
+            workflow_id="mock-workflow-id",
+        )
+
+        constructor = DurableContext if self.is_durable else Context
+
+        return constructor(
+            action=action,
+            dispatcher_client=self.workflow.client._client.dispatcher,
+            admin_client=self.workflow.client._client.admin,
+            event_client=self.workflow.client._client.event,
+            durable_event_listener=None,
+            worker=WorkerContext(
+                labels={}, client=self.workflow.client._client.dispatcher
+            ),
+            runs_client=self.workflow.client._client.runs,
+            lifespan_context=lifespan_context,
+            log_sender=AsyncLogSender(self.workflow.client._client.event),
+        )
+
+    def mock_run(
+        self,
+        input: TWorkflowInput | None = None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan: Any = None,
+    ) -> R:
+        """
+        Mimic the execution of a task. This method is intended to be used to unit test
+        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
+        tasks and `aio_mock_run` for async tasks.
+
+        :param input: The input to the task.
+        :param additional_metadata: Additional metadata to attach to the task.
+        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
+        :param retry_count: The number of times the task has been retried.
+        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
+
+        :return: The output of the task.
+        :raises TypeError: If the task is an async function and `mock_run` is called, or if the task is a sync function and `aio_mock_run` is called.
+        """
+
+        if self.is_async_function:
+            raise TypeError(
+                f"{self.name} is not a sync function. Use `aio_mock_run` instead."
+            )
+
+        ctx = self._create_mock_context(
+            input, additional_metadata, parent_outputs, retry_count, lifespan
+        )
+
+        return self.call(ctx)
+
+    async def aio_mock_run(
+        self,
+        input: TWorkflowInput | None = None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan: Any = None,
+    ) -> R:
+        """
+        Mimic the execution of a task. This method is intended to be used to unit test
+        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
+        tasks and `aio_mock_run` for async tasks.
+
+        :param input: The input to the task.
+        :param additional_metadata: Additional metadata to attach to the task.
+        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
+        :param retry_count: The number of times the task has been retried.
+        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
+
+        :return: The output of the task.
+        :raises TypeError: If the task is an async function and `mock_run` is called, or if the task is a sync function and `aio_mock_run` is called.
+        """
+
+        if not self.is_async_function:
+            raise TypeError(
+                f"{self.name} is not an async function. Use `mock_run` instead."
+            )
+
+        ctx = self._create_mock_context(
+            input,
+            additional_metadata,
+            parent_outputs,
+            retry_count,
+            lifespan,
+        )
+
+        return await self.aio_call(ctx)

hatchet_sdk/runnables/workflow.py

@@ -2,7 +2,16 @@ import asyncio
 from collections.abc import Callable
 from datetime import datetime, timedelta
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_type_hints
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Literal,
+    TypeVar,
+    cast,
+    get_type_hints,
+    overload,
+)
 
 from google.protobuf import timestamp_pb2
 from pydantic import BaseModel, model_validator
@@ -651,39 +660,83 @@ class Workflow(BaseWorkflow[TWorkflowInput]):
 
         return await ref.aio_result()
 
+    def _get_result(
+        self, ref: WorkflowRunRef, return_exceptions: bool
+    ) -> dict[str, Any] | BaseException:
+        try:
+            return ref.result()
+        except Exception as e:
+            if return_exceptions:
+                return e
+            raise e
+
+    @overload
+    def run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: Literal[True],
+    ) -> list[dict[str, Any] | BaseException]: ...
+
+    @overload
     def run_many(
         self,
         workflows: list[WorkflowRunTriggerConfig],
-    ) -> list[dict[str, Any]]:
+        return_exceptions: Literal[False] = False,
+    ) -> list[dict[str, Any]]: ...
+
+    def run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: bool = False,
+    ) -> list[dict[str, Any]] | list[dict[str, Any] | BaseException]:
         """
         Run a workflow in bulk and wait for all runs to complete.
         This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.
 
         :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
+        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
         :returns: A list of results for each workflow run.
         """
         refs = self.client._client.admin.run_workflows(
             workflows=workflows,
         )
 
-        return [ref.result() for ref in refs]
+        return [self._get_result(ref, return_exceptions) for ref in refs]
 
+    @overload
     async def aio_run_many(
         self,
         workflows: list[WorkflowRunTriggerConfig],
-    ) -> list[dict[str, Any]]:
+        return_exceptions: Literal[True],
+    ) -> list[dict[str, Any] | BaseException]: ...
+
+    @overload
+    async def aio_run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: Literal[False] = False,
+    ) -> list[dict[str, Any]]: ...
+
+    async def aio_run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: bool = False,
+    ) -> list[dict[str, Any]] | list[dict[str, Any] | BaseException]:
         """
         Run a workflow in bulk and wait for all runs to complete.
         This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.
 
         :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
+        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
         :returns: A list of results for each workflow run.
         """
         refs = await self.client._client.admin.aio_run_workflows(
             workflows=workflows,
         )
 
-        return await asyncio.gather(*[ref.aio_result() for ref in refs])
+        return await asyncio.gather(
+            *[ref.aio_result() for ref in refs], return_exceptions=return_exceptions
+        )
 
     def run_many_no_wait(
         self,
@@ -946,7 +999,7 @@ class Workflow(BaseWorkflow[TWorkflowInput]):
 
         :param backoff_max_seconds: The maximum number of seconds to allow retries with exponential backoff to continue.
 
-        :param concurrency: A list of concurrency expressions for the on-success task.
+        :param concurrency: A list of concurrency expressions for the on-failure task.
 
         :returns: A decorator which creates a `Task` object.
         """
@@ -1137,7 +1190,18 @@ class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
 
         self.config = self._workflow.config
 
-    def _extract_result(self, result: dict[str, Any]) -> R:
+    @overload
+    def _extract_result(self, result: dict[str, Any]) -> R: ...
+
+    @overload
+    def _extract_result(self, result: BaseException) -> BaseException: ...
+
+    def _extract_result(
+        self, result: dict[str, Any] | BaseException
+    ) -> R | BaseException:
+        if isinstance(result, BaseException):
+            return result
+
         output = result.get(self._task.name)
 
         if not self._output_validator:
@@ -1217,30 +1281,72 @@ class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
 
         return TaskRunRef[TWorkflowInput, R](self, ref)
 
-    def run_many(self, workflows: list[WorkflowRunTriggerConfig]) -> list[R]:
+    @overload
+    def run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: Literal[True],
+    ) -> list[R | BaseException]: ...
+
+    @overload
+    def run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: Literal[False] = False,
+    ) -> list[R]: ...
+
+    def run_many(
+        self, workflows: list[WorkflowRunTriggerConfig], return_exceptions: bool = False
+    ) -> list[R] | list[R | BaseException]:
         """
         Run a workflow in bulk and wait for all runs to complete.
         This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.
 
         :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
+        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
         :returns: A list of results for each workflow run.
         """
         return [
             self._extract_result(result)
-            for result in self._workflow.run_many(workflows)
+            for result in self._workflow.run_many(
+                workflows,
+                ## hack: typing needs literal
+                True if return_exceptions else False,  # noqa: SIM210
+            )
         ]
 
-    async def aio_run_many(self, workflows: list[WorkflowRunTriggerConfig]) -> list[R]:
+    @overload
+    async def aio_run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: Literal[True],
+    ) -> list[R | BaseException]: ...
+
+    @overload
+    async def aio_run_many(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+        return_exceptions: Literal[False] = False,
+    ) -> list[R]: ...
+
+    async def aio_run_many(
+        self, workflows: list[WorkflowRunTriggerConfig], return_exceptions: bool = False
+    ) -> list[R] | list[R | BaseException]:
         """
         Run a workflow in bulk and wait for all runs to complete.
         This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.
 
         :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
+        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
         :returns: A list of results for each workflow run.
         """
         return [
             self._extract_result(result)
-            for result in await self._workflow.aio_run_many(workflows)
+            for result in await self._workflow.aio_run_many(
+                workflows,
+                ## hack: typing needs literal
+                True if return_exceptions else False,  # noqa: SIM210
+            )
         ]
 
     def run_many_no_wait(
@@ -1273,3 +1379,104 @@ class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
         refs = await self._workflow.aio_run_many_no_wait(workflows)
 
         return [TaskRunRef[TWorkflowInput, R](self, ref) for ref in refs]
+
+    def mock_run(
+        self,
+        input: TWorkflowInput | None = None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan: Any = None,
+    ) -> R:
+        """
+        Mimic the execution of a task. This method is intended to be used to unit test
+        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
+        tasks and `aio_mock_run` for async tasks.
+
+        :param input: The input to the task.
+        :param additional_metadata: Additional metadata to attach to the task.
+        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
+        :param retry_count: The number of times the task has been retried.
+        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
+
+        :return: The output of the task.
+        """
+
+        return self._task.mock_run(
+            input=input,
+            additional_metadata=additional_metadata,
+            parent_outputs=parent_outputs,
+            retry_count=retry_count,
+            lifespan=lifespan,
+        )
+
+    async def aio_mock_run(
+        self,
+        input: TWorkflowInput | None = None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan: Any = None,
+    ) -> R:
+        """
+        Mimic the execution of a task. This method is intended to be used to unit test
+        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
+        tasks and `aio_mock_run` for async tasks.
+
+        :param input: The input to the task.
+        :param additional_metadata: Additional metadata to attach to the task.
+        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
+        :param retry_count: The number of times the task has been retried.
+        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
+
+        :return: The output of the task.
+        """
+
+        return await self._task.aio_mock_run(
+            input=input,
+            additional_metadata=additional_metadata,
+            parent_outputs=parent_outputs,
+            retry_count=retry_count,
+            lifespan=lifespan,
+        )
+
+    @property
+    def is_async_function(self) -> bool:
+        """
+        Check if the task is an async function.
+
+        :returns: True if the task is an async function, False otherwise.
+        """
+        return self._task.is_async_function
+
+    def get_run_ref(self, run_id: str) -> TaskRunRef[TWorkflowInput, R]:
+        """
+        Get a reference to a task run by its run ID.
+
+        :param run_id: The ID of the run to get the reference for.
+        :returns: A `TaskRunRef` object representing the reference to the task run.
+        """
+        wrr = self._workflow.client._client.runs.get_run_ref(run_id)
+        return TaskRunRef[TWorkflowInput, R](self, wrr)
+
+    async def aio_get_result(self, run_id: str) -> R:
+        """
+        Get the result of a task run by its run ID.
+
+        :param run_id: The ID of the run to get the result for.
+        :returns: The result of the task run.
+        """
+        run_ref = self.get_run_ref(run_id)
+
+        return await run_ref.aio_result()
+
+    def get_result(self, run_id: str) -> R:
+        """
+        Get the result of a task run by its run ID.
+
+        :param run_id: The ID of the run to get the result for.
+        :returns: The result of the task run.
+        """
+        run_ref = self.get_run_ref(run_id)
+
+        return run_ref.result()