Cache middleware properly

dramatiq_workflow/_base.py

@@ -164,18 +164,18 @@ def __augment_message(self, message: Message, completion_callbacks: SerializedCo
 
     @property
     def __middleware(self) -> WorkflowMiddleware:
-        if not hasattr(self, "__cached_middleware"):
+        if not hasattr(self, "_cached_middleware"):
             for middleware in self.broker.middleware:
                 if isinstance(middleware, WorkflowMiddleware):
-                    self.__cached_middleware = middleware
+                    self._cached_middleware = middleware
                     break
             else:
                 raise RuntimeError(
                     "WorkflowMiddleware middleware not found! Did you forget "
                     "to set it up? It is required if you want to use "
                     "workflows."
                 )
-        return self.__cached_middleware
+        return self._cached_middleware
 
     @property
     def __rate_limiter_backend(self) -> dramatiq.rate_limits.RateLimiterBackend:

dramatiq_workflow/tests/test_workflow.py

@@ -12,14 +12,18 @@ class WorkflowTests(unittest.TestCase):
     def setUp(self):
         self.rate_limiter_backend = mock.create_autospec(dramatiq.rate_limits.RateLimiterBackend, instance=True)
         self.barrier = mock.create_autospec(dramatiq.rate_limits.Barrier)
-        self.broker = mock.MagicMock(
-            middleware=[
-                WorkflowMiddleware(
-                    rate_limiter_backend=self.rate_limiter_backend,
-                    barrier_type=self.barrier,
-                )
+        self.broker = mock.MagicMock()
+        self.workflow_middleware = WorkflowMiddleware(
+            rate_limiter_backend=self.rate_limiter_backend,
+            barrier_type=self.barrier,
+        )
+        self.middleware_list = mock.PropertyMock(
+            return_value=[
+                self.workflow_middleware,
             ]
         )
+        type(self.broker).middleware = self.middleware_list
+
         self.task = mock.MagicMock()
         self.task.message.side_effect = lambda *args, **kwargs: self.__make_message(
             self.__generate_id(), *args, **kwargs
@@ -389,3 +393,13 @@ def test_nested_delays(self, time_mock):
             ),
             delay=20,
         )
+
+    def test_middleware_is_cached(self):
+        workflow = Workflow(Chain(self.task.message(), self.task.message()), broker=self.broker)
+
+        # Access middleware property multiple times
+        workflow.run()
+        workflow.run()
+
+        # Check that broker.middleware was accessed only once
+        self.middleware_list.assert_called_once()