test: use real stream subscriptions for workflow dispatch (#43213)

api/tests/unit_tests/services/test_installed_app_generation_adapters.py

@@ -8,6 +8,7 @@ from uuid import uuid4
 
 import pytest
 from flask import Flask
+from pytest_mock import MockerFixture
 from sqlalchemy import event, inspect, update
 from sqlalchemy.orm import Session, sessionmaker
 
@@ -17,7 +18,12 @@ from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.features.rate_limiting.rate_limit import RateLimit, RateLimitGenerator
 from enums import DeploymentEdition
 from extensions.ext_database import db
-from libs.broadcast_channel.channel import BroadcastChannel, Subscription, SupportsPreparedSubscription, Topic
+from extensions.ext_redis import RedisClientWrapper
+from libs.broadcast_channel.redis.streams_channel import (
+    StreamsBroadcastChannel,
+    _StreamsSubscriber,
+    _StreamsSubscription,
+)
 from models import Account, App, AppMode, AppModelConfig, Conversation, Message, Workflow
 from models.enums import ConversationFromSource
 from models.workflow import WorkflowType
@@ -701,6 +707,8 @@ def test_workflow_dispatch_starts_task_after_subscription_with_runtime_session_c
     config_overrides: Callable[..., None],
     mode: AppMode,
     trigger: bool,
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
+    mocker: MockerFixture,
 ) -> None:
     config_overrides(
         DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY,
@@ -737,22 +745,24 @@ def test_workflow_dispatch_starts_task_after_subscription_with_runtime_session_c
         workflow_id, tenant_id = workflow.id, app.tenant_id
 
     transport_events: list[str] = []
-    subscription = MagicMock(spec=Subscription)
-    subscriber = MagicMock(spec=SupportsPreparedSubscription)
-    subscriber.prepare_subscription.return_value = subscription
-    topic = MagicMock(spec=Topic)
-    topic.as_subscriber.return_value = subscriber
-    channel = MagicMock(spec=BroadcastChannel)
-    channel.topic.return_value = topic
+    redis, commands = redis_transport
+    empty_stream: list[tuple[bytes, dict[bytes, bytes]]] = []
+    commands.return_value = empty_stream
+    channel = StreamsBroadcastChannel(redis._require_client())
+    prepare = mocker.spy(_StreamsSubscriber, "prepare_subscription")
+    close = mocker.spy(_StreamsSubscription, "close")
 
-    def activate_subscription() -> Subscription:
+    def activate_subscription(subscription: _StreamsSubscription) -> None:
         assert len(harness.closed_sessions) == 2
         assert all(not session.in_transaction() for session in harness.closed_sessions)
-        transport_events.append("subscribe")
-        return subscription
+        if not transport_events:
+            transport_events.append("subscribe")
+            subscription._queue.put_nowait(b'{"event":"workflow_finished"}')
 
-    subscription.__enter__.side_effect = activate_subscription
-    subscription.receive.return_value = b'{"event":"workflow_finished"}'
+    # Keep the real preparation, queue and lifecycle; replace only listener-thread startup.
+    start_listener = mocker.patch.object(
+        _StreamsSubscription, "_start_if_needed", autospec=True, side_effect=activate_subscription
+    )
     monkeypatch.setattr(message_based_app_generator, "get_pubsub_broadcast_channel", lambda: channel)
     submitted: list[generation_module.AppExecutionParams] = []
 
@@ -771,8 +781,12 @@ def test_workflow_dispatch_starts_task_after_subscription_with_runtime_session_c
     assert len(harness.closed_sessions) == 2
     assert harness.committed_sessions == [harness.closed_sessions[1]]
     assert submitted == []
-    subscriber.prepare_subscription.assert_called_once_with()
-    subscription.__enter__.assert_not_called()
+    prepare.assert_called_once()
+    subscription = prepare.spy_return
+    assert isinstance(subscription, _StreamsSubscription)
+    assert subscription._start_id == "0-0"
+    start_listener.assert_not_called()
+    assert commands.call_args.args[0] == "XREVRANGE"
     assert next(result) == "event: ping\n\n"
     assert submitted == []
     assert json.loads(next(result).removeprefix("data: ")) == {"event": "workflow_finished"}
@@ -786,12 +800,15 @@ def test_workflow_dispatch_starts_task_after_subscription_with_runtime_session_c
     assert payload.args == args
     assert payload.invoke_from == InvokeFrom.EXPLORE
     assert payload.streaming is True
-    subscription.__exit__.assert_called_once()
+    close.assert_called_once_with(subscription)
+    assert subscription._closed is True
     assert result.closed is True
 
 
 def test_unpublished_workflow_raises_before_subscription_or_task_creation(
     harness: _RuntimeHarness,
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
+    mocker: MockerFixture,
     monkeypatch: pytest.MonkeyPatch,
     sqlite_session_factory: sessionmaker[Session],
     config_overrides: Callable[..., None],
@@ -818,7 +835,9 @@ def test_unpublished_workflow_raises_before_subscription_or_task_creation(
             )
         )
 
-    channel = MagicMock(spec=BroadcastChannel)
+    redis, commands = redis_transport
+    channel = StreamsBroadcastChannel(redis._require_client())
+    topic = mocker.spy(channel, "topic")
     monkeypatch.setattr(message_based_app_generator, "get_pubsub_broadcast_channel", lambda: channel)
     enqueue = MagicMock()
     monkeypatch.setattr(generation_module.workflow_based_app_execution_task, "delay", enqueue)
@@ -830,5 +849,6 @@ def test_unpublished_workflow_raises_before_subscription_or_task_creation(
     assert len(harness.closed_sessions) == 2
     assert all(not session.in_transaction() for session in harness.closed_sessions)
     assert harness.committed_sessions == []
-    channel.topic.assert_not_called()
+    topic.assert_not_called()
+    commands.assert_not_called()
     enqueue.assert_not_called()