From c294006ecfa2beae6558a8e314379dcc8328771f Mon Sep 17 00:00:00 2001
From: "yunlu.wen" <yunlu.wen@dify.ai>
Date: Fri, 17 Apr 2026 11:10:11 +0800
Subject: [PATCH] fix unit tests

---
 api/tests/unit_tests/core/app/apps/test_streaming_utils.py | 7 ++++---
 .../test_workflow_event_snapshot_service_additional.py     | 7 ++++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/api/tests/unit_tests/core/app/apps/test_streaming_utils.py b/api/tests/unit_tests/core/app/apps/test_streaming_utils.py
index a7714c56ce..04bd651e52 100644
--- a/api/tests/unit_tests/core/app/apps/test_streaming_utils.py
+++ b/api/tests/unit_tests/core/app/apps/test_streaming_utils.py
@@ -12,10 +12,11 @@ from models.model import AppMode
 
 
 class FakeSubscription:
-    def __init__(self, message_queue: queue.Queue[bytes], state: dict[str, bool]) -> None:
+    def __init__(self, message_queue: queue.Queue[bytes], state: dict[str, bool], replay: bool = False) -> None:
         self._queue = message_queue
         self._state = state
         self._closed = False
+        self._replay = replay
 
     def __enter__(self):
         self._state["subscribed"] = True
@@ -43,8 +44,8 @@ class FakeTopic:
         self._queue: queue.Queue[bytes] = queue.Queue()
         self._state = {"subscribed": False}
 
-    def subscribe(self) -> FakeSubscription:
-        return FakeSubscription(self._queue, self._state)
+    def subscribe(self, replay: bool = False) -> FakeSubscription:
+        return FakeSubscription(self._queue, self._state, replay)
 
     def publish(self, payload: bytes) -> None:
         self._queue.put(payload)
diff --git a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py
index d2634d7d7b..ad102d54dd 100644
--- a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py
+++ b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py
@@ -96,8 +96,9 @@ class _SessionMaker:
 
 
 class _SubscriptionContext:
-    def __init__(self, subscription: Any) -> None:
+    def __init__(self, subscription: Any, replay: bool = False) -> None:
         self._subscription = subscription
+        self._replay = replay
 
     def __enter__(self) -> Any:
         return self._subscription
@@ -110,8 +111,8 @@ class _Topic:
     def __init__(self, subscription: Any) -> None:
         self._subscription = subscription
 
-    def subscribe(self) -> _SubscriptionContext:
-        return _SubscriptionContext(self._subscription)
+    def subscribe(self, replay: bool = False) -> _SubscriptionContext:
+        return _SubscriptionContext(self._subscription, replay)
 
 
 class _StaticSubscription: