Add change supporting unit testing

azure/durable_functions/decorators/durable_app.py

@@ -198,6 +198,17 @@ async def df_client_middleware(*args, **kwargs):
             # Invoke user code with rich DF Client binding
             return await user_code(*args, **kwargs)
 
+        # Todo: This feels awkward - however, there are two reasons that I can't naively implement
+        # this in the same way as entities and orchestrators:
+        # 1. We intentionally wrap this exported signature with @wraps, to preserve the original
+        #    signature of the user code. This means that we can't just assign a new object to the
+        #    fb._function._func, as that would overwrite the original signature.
+        # 2. I have not yet fully tested the behavior of overriding __call__ on an object with an
+        #    async method.
+        # Here we lose type hinting and auto-documentation - not great. Need to find a better way
+        # to do this. 
+        df_client_middleware.client_function = fb._function._func
+
         user_code_with_rich_client = df_client_middleware
         fb._function._func = user_code_with_rich_client
 

azure/durable_functions/entity.py

@@ -3,12 +3,45 @@
 from datetime import datetime
 from typing import Callable, Any, List, Dict
 
+import azure.functions as func
 
 class InternalEntityException(Exception):
     """Framework-internal Exception class (for internal use only)."""
 
     pass
 
+class EntityHandler(Callable):
+    """Durable Entity Handler.
+    A callable class that wraps the user defined entity function for execution by the Python worker
+    and also allows access to the original method for unit testing
+    """
+
+    def __init__(self, func: Callable[[DurableEntityContext], None]):
+        """
+        Create a new entity handler for the user defined entity function.
+
+        Parameters
+        ----------
+        func: Callable[[DurableEntityContext], None]
+            The user defined entity function.
+        """
+        self.entity_function = func
+
+    def __call__(self, context: func.EntityContext) -> str:
+        """Handle the execution of the user defined entity function.
+        Parameters
+        ----------
+        context : func.EntityContext
+            The DF entity context"""
+        # It is not clear when the context JSON would be found
+        # inside a "body"-key, but this pattern matches the
+        # orchestrator implementation, so we keep it for safety.
+        context_body = getattr(context, "body", None)
+        if context_body is None:
+            context_body = context
+        ctx, batch = DurableEntityContext.from_json(context_body)
+        return Entity(self.entity_function).handle(ctx, batch)
+
 
 class Entity:
     """Durable Entity Class.
@@ -92,19 +125,10 @@ def create(cls, fn: Callable[[DurableEntityContext], None]) -> Callable[[Any], s
 
         Returns
         -------
-        Callable[[Any], str]
-            Handle function of the newly created entity client
+        EntityHandler
+            Entity Handler callable for the newly created entity client
         """
-        def handle(context) -> str:
-            # It is not clear when the context JSON would be found
-            # inside a "body"-key, but this pattern matches the
-            # orchestrator implementation, so we keep it for safety.
-            context_body = getattr(context, "body", None)
-            if context_body is None:
-                context_body = context
-            ctx, batch = DurableEntityContext.from_json(context_body)
-            return Entity(fn).handle(ctx, batch)
-        return handle
+        return EntityHandler(fn)
 
     def _elapsed_milliseconds_since(self, start_time: datetime) -> int:
         """Calculate the elapsed time, in milliseconds, from the start_time to the present.

azure/durable_functions/orchestrator.py

@@ -11,6 +11,35 @@
 import azure.functions as func
 
 
+class OrchestrationHandler(Callable):
+    """Durable Orchestration Handler.
+    A callable class that wraps the user defined generator function for execution by the Python worker
+    and also allows access to the original method for unit testing
+    """
+
+    def __init__(self, func: Callable[[DurableOrchestrationContext], Generator[Any, Any, Any]]):
+        """
+        Create a new orchestrator handler for the user defined orchestrator function.
+
+        Parameters
+        ----------
+        func: Callable[[DurableOrchestrationContext], Generator[Any, Any, Any]]
+            The user defined orchestrator function.
+        """
+        self.orchestrator_function = func
+
+    def __call__(self, context: func.OrchestrationContext) -> str:
+        """Handle the execution of the user defined orchestrator function.
+        Parameters
+        ----------
+        context : func.OrchestrationContext
+            The DF orchestration context"""
+        context_body = getattr(context, "body", None)
+        if context_body is None:
+            context_body = context
+        return Orchestrator(self.orchestrator_function).handle(DurableOrchestrationContext.from_json(context_body))
+
+
 class Orchestrator:
     """Durable Orchestration Class.
 
@@ -58,14 +87,7 @@ def create(cls, fn: Callable[[DurableOrchestrationContext], Generator[Any, Any,
 
         Returns
         -------
-        Callable[[Any], str]
-            Handle function of the newly created orchestration client
+        OrchestrationHandler
+            Orchestration handler callable class for the newly created orchestration client
         """
-
-        def handle(context: func.OrchestrationContext) -> str:
-            context_body = getattr(context, "body", None)
-            if context_body is None:
-                context_body = context
-            return Orchestrator(fn).handle(DurableOrchestrationContext.from_json(context_body))
-
-        return handle
+        return OrchestrationHandler(fn)