Update unit testing for compatibility with worker

azure/durable_functions/entity.py

@@ -3,52 +3,13 @@
 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.
 
@@ -131,10 +92,22 @@ def create(cls, fn: Callable[[DurableEntityContext], None]) -> Callable[[Any], s
 
         Returns
         -------
-        EntityHandler
-            Entity Handler callable for the newly created entity client
+        Callable[[Any], str]
+            Handle function of the newly created entity client
         """
-        return EntityHandler(fn)
+        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)
+
+        handle.entity_function = fn
+
+        return handle
 
     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,41 +11,6 @@
 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.
 
@@ -93,7 +58,16 @@ def create(cls, fn: Callable[[DurableOrchestrationContext], Generator[Any, Any,
 
         Returns
         -------
-        OrchestrationHandler
-            Orchestration handler callable class for the newly created orchestration client
+        Callable[[Any], str]
+            Handle function of the newly created orchestration client
         """
-        return OrchestrationHandler(fn)
+
+        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))
+
+        handle.orchestrator_function = fn
+
+        return handle

samples-v2/function_chaining/tests/test_my_orchestrator.py

@@ -31,7 +31,7 @@ def test_chaining_orchestrator(self, context):
     expected_activity_calls = [call('say_hello', 'Tokyo'),
                                call('say_hello', 'Seattle'),
                                call('say_hello', 'London')]
-    
+
     self.assertEqual(context.call_activity.call_count, 3)
     self.assertEqual(context.call_activity.call_args_list, expected_activity_calls)
     self.assertEqual(values[3], ["Hello Tokyo!", "Hello Seattle!", "Hello London!"])