Use functools.wraps on when wrapping callbacks (#137)

This preserves the wrapped callback name when stacking multiple
middleware layers

Add subscribe_broadcast and subscribe_temporary to TimerMiddleware

Author: rjgildea

src/workflows/transport/middleware/__init__.py

@@ -178,6 +178,45 @@ def __init__(self, logger: logging.Logger = None, level=logging.INFO):
         self.level = level
 
     def subscribe(self, call_next: Callable, channel, callback, **kwargs) -> int:
+        @functools.wraps(callback)
+        def wrapped_callback(header, message):
+            start_time = time.perf_counter()
+            result = callback(header, message)
+            end_time = time.perf_counter()
+            source = get_callback_source(callback)
+            self.logger.log(
+                self.level,
+                f"Callback for {source} took {end_time - start_time:.4f} seconds",
+            )
+            return result
+
+        return call_next(channel, wrapped_callback, **kwargs)
+
+    def subscribe_temporary(
+        self,
+        call_next: Callable,
+        channel_hint: Optional[str],
+        callback: MessageCallback,
+        **kwargs,
+    ) -> TemporarySubscription:
+        @functools.wraps(callback)
+        def wrapped_callback(header, message):
+            start_time = time.perf_counter()
+            result = callback(header, message)
+            end_time = time.perf_counter()
+            source = get_callback_source(callback)
+            self.logger.log(
+                self.level,
+                f"Callback for {source} took {end_time - start_time:.4f} seconds",
+            )
+            return result
+
+        return call_next(channel_hint, wrapped_callback, **kwargs)
+
+    def subscribe_broadcast(
+        self, call_next: Callable, channel, callback, **kwargs
+    ) -> int:
+        @functools.wraps(callback)
         def wrapped_callback(header, message):
             start_time = time.perf_counter()
             result = callback(header, message)

src/workflows/transport/middleware/prometheus.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import time
 from typing import Callable, Optional
 
@@ -82,6 +83,7 @@ def __init__(self, source: str):
         self.source = source
 
     def subscribe(self, call_next: Callable, channel, callback, **kwargs) -> int:
+        @functools.wraps(callback)
         def wrapped_callback(header, message):
             start_time = time.perf_counter()
             result = callback(header, message)
@@ -102,6 +104,7 @@ def subscribe_temporary(
         callback: MessageCallback,
         **kwargs,
     ) -> TemporarySubscription:
+        @functools.wraps(callback)
         def wrapped_callback(header, message):
             start_time = time.perf_counter()
             result = callback(header, message)
@@ -118,6 +121,7 @@ def wrapped_callback(header, message):
     def subscribe_broadcast(
         self, call_next: Callable, channel, callback, **kwargs
     ) -> int:
+        @functools.wraps(callback)
         def wrapped_callback(header, message):
             start_time = time.perf_counter()
             result = callback(header, message)

tests/transport/test_middleware.py

@@ -54,15 +54,34 @@ def test_timer_middleware(caplog):
     def callback(header, message):
         time.sleep(1)
 
-    subscription_id = offline.subscribe(str(mock.sentinel.channel), callback)
+    sid_1 = offline.subscribe(str(mock.sentinel.channel), callback)
+    sid_2 = offline.subscribe_broadcast(str(mock.sentinel.channel), callback)
+    ts = offline.subscribe_temporary(str(mock.sentinel.channel), callback)
+
+    expected_text = (
+        "Callback for test_middleware:test_timer_middleware.<locals>.callback took"
+    )
+
     with caplog.at_level(logging.DEBUG):
-        offline.subscription_callback(subscription_id)(
+        offline.subscription_callback(sid_1)(
             {"destination": "foo"}, str(mock.sentinel.message)
         )
-        assert (
-            "Callback for test_middleware:test_timer_middleware.<locals>.callback took"
-            in caplog.text
+        assert expected_text in caplog.text
+        caplog.clear()
+
+    with caplog.at_level(logging.DEBUG):
+        offline.subscription_callback(sid_2)(
+            {"destination": "bar"}, str(mock.sentinel.message)
+        )
+        assert expected_text in caplog.text
+        caplog.clear()
+
+    with caplog.at_level(logging.DEBUG):
+        offline.subscription_callback(ts.subscription_id)(
+            {"destination": "foobar"}, str(mock.sentinel.message)
         )
+        assert expected_text in caplog.text
+        caplog.clear()
 
 
 def test_prometheus_middleware():