test: exercise collaboration leader and session storage (#43230)

api/tests/unit_tests/services/test_workflow_collaboration_service.py

@@ -1,6 +1,6 @@
 import logging
 from collections.abc import Iterator
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
@@ -63,13 +63,6 @@ def _app(*, app_id: str, tenant_id: str, maintainer: str | None = None) -> App:
 
 
 class TestWorkflowCollaborationService:
-    @pytest.fixture
-    def service(self) -> tuple[WorkflowCollaborationService, Mock, Mock]:
-        repository = Mock(spec=WorkflowCollaborationRepository)
-        repository.graph_view_state_lock.return_value = nullcontext()
-        socketio = Mock()
-        return WorkflowCollaborationService(repository, socketio, server_id="server-1"), repository, socketio
-
     def test_authorize_and_join_workflow_room_returns_leader_status(
         self, real_service: ServiceFixture, db_session: Session
     ) -> None:
@@ -996,12 +989,10 @@ class TestWorkflowCollaborationService:
         handle_leader_disconnect.assert_called_once_with("wf-1", "sid-1")
         broadcast_online_users.assert_called_once_with("wf-1")
 
-    def test_get_or_set_leader_returns_active_leader(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
+    def test_get_or_set_leader_returns_active_leader(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = "sid-1"
+        collaboration_service, repository, _socketio = real_service
+        repository.set_leader("wf-1", "sid-1")
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             # Act
@@ -1009,25 +1000,25 @@ class TestWorkflowCollaborationService:
 
         # Assert
         assert result == "sid-1"
-        repository.set_leader_if_absent.assert_not_called()
+        assert repository.get_current_leader("wf-1") == "sid-1"
 
-    def test_get_or_set_leader_replaces_dead_leader(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
+    def test_get_or_set_leader_replaces_dead_leader(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = "sid-1"
-        repository.set_leader_if_absent.return_value = True
-        repository.list_sessions.return_value = [
-            {
-                "user_id": "u-2",
-                "username": "B",
-                "avatar": None,
-                "sid": "sid-2",
-                "connected_at": 1,
-                "graph_active": True,
-            }
-        ]
+        collaboration_service, repository, _socketio = real_service
+        repository.set_leader("wf-1", "sid-1")
+        seed_sessions(
+            repository,
+            [
+                {
+                    "user_id": "u-2",
+                    "username": "B",
+                    "avatar": None,
+                    "sid": "sid-2",
+                    "connected_at": 1,
+                    "graph_active": True,
+                }
+            ],
+        )
 
         with (
             patch.object(collaboration_service, "is_session_active", side_effect=lambda _wf, sid: sid != "sid-1"),
@@ -1038,61 +1029,73 @@ class TestWorkflowCollaborationService:
 
         # Assert
         assert result == "sid-2"
-        repository.delete_session.assert_called_once_with("wf-1", "sid-1")
-        repository.delete_leader.assert_called_once_with("wf-1")
+        assert not repository.session_exists("wf-1", "sid-1")
         broadcast_leader_change.assert_called_once_with("wf-1", "sid-2")
 
-    def test_get_or_set_leader_falls_back_to_existing(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
+        assert repository.get_current_leader("wf-1") == "sid-2"
+
+    def test_get_or_set_leader_falls_back_to_existing(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.side_effect = [None, "sid-3"]
-        repository.set_leader_if_absent.return_value = False
-        repository.list_sessions.return_value = [
-            {
-                "user_id": "u-2",
-                "username": "B",
-                "avatar": None,
-                "sid": "sid-2",
-                "connected_at": 1,
-                "graph_active": True,
-            }
-        ]
+        collaboration_service, repository, _socketio = real_service
+        set_if_absent = repository.set_leader_if_absent
+
+        def competing_leader(workflow_id: str, sid: str) -> bool:
+            repository.set_leader(workflow_id, "sid-3")
+            return set_if_absent(workflow_id, sid)
+
+        seed_sessions(
+            repository,
+            [
+                {
+                    "user_id": "u-2",
+                    "username": "B",
+                    "avatar": None,
+                    "sid": "sid-2",
+                    "connected_at": 1,
+                    "graph_active": True,
+                }
+            ],
+        )
 
         # Act
-        result = collaboration_service.get_or_set_leader("wf-1", "sid-2")
+        with patch.object(repository, "set_leader_if_absent", side_effect=competing_leader):
+            result = collaboration_service.get_or_set_leader("wf-1", "sid-2")
 
         # Assert
         assert result == "sid-3"
 
-    def test_get_or_set_leader_returns_sid_when_leader_still_missing(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.side_effect = [None, None]
-        repository.set_leader_if_absent.return_value = False
+    def test_get_or_set_leader_returns_sid_when_leader_still_missing(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        set_if_absent = repository.set_leader_if_absent
 
-        result = collaboration_service.get_or_set_leader("wf-1", "sid-2")
+        def disappearing_leader(workflow_id: str, sid: str) -> bool:
+            repository.set_leader(workflow_id, "sid-3")
+            acquired = set_if_absent(workflow_id, sid)
+            repository.delete_leader(workflow_id)
+            return acquired
+
+        with patch.object(repository, "set_leader_if_absent", side_effect=disappearing_leader):
+            result = collaboration_service.get_or_set_leader("wf-1", "sid-2")
 
         assert result == "sid-2"
 
-    def test_handle_leader_disconnect_elects_new(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
+    def test_handle_leader_disconnect_elects_new(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = "sid-1"
-        repository.list_sessions.return_value = [
-            {
-                "user_id": "u-2",
-                "username": "B",
-                "avatar": None,
-                "sid": "sid-2",
-                "connected_at": 1,
-                "graph_active": True,
-            }
-        ]
+        collaboration_service, repository, _socketio = real_service
+        repository.set_leader("wf-1", "sid-1")
+        seed_sessions(
+            repository,
+            [
+                {
+                    "user_id": "u-2",
+                    "username": "B",
+                    "avatar": None,
+                    "sid": "sid-2",
+                    "connected_at": 1,
+                    "graph_active": True,
+                }
+            ],
+        )
 
         with (
             patch.object(collaboration_service, "is_session_active", return_value=True),
@@ -1102,44 +1105,41 @@ class TestWorkflowCollaborationService:
             collaboration_service.handle_leader_disconnect("wf-1", "sid-1")
 
         # Assert
-        repository.set_leader.assert_called_once_with("wf-1", "sid-2")
+        assert repository.get_current_leader("wf-1") == "sid-2"
         broadcast_leader_change.assert_called_once_with("wf-1", "sid-2")
 
-    def test_handle_leader_disconnect_clears_when_empty(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
+    def test_handle_leader_disconnect_clears_when_empty(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = "sid-1"
-        repository.list_sessions.return_value = []
+        collaboration_service, repository, _socketio = real_service
+        repository.set_leader("wf-1", "sid-1")
+        seed_sessions(repository, [])
 
         # Act
         collaboration_service.handle_leader_disconnect("wf-1", "sid-1")
 
         # Assert
-        repository.delete_leader.assert_called_once_with("wf-1")
 
-    def test_handle_leader_disconnect_ignores_non_leader_or_missing_leader(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
+        assert repository.get_current_leader("wf-1") is None
 
-        repository.get_current_leader.return_value = None
+    def test_handle_leader_disconnect_ignores_non_leader_or_missing_leader(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+
+        assert repository.get_current_leader("wf-1") is None
         collaboration_service.handle_leader_disconnect("wf-1", "sid-1")
 
-        repository.get_current_leader.return_value = "sid-leader"
+        repository.set_leader("wf-1", "sid-leader")
         collaboration_service.handle_leader_disconnect("wf-1", "sid-other")
 
-        repository.set_leader.assert_not_called()
-        repository.delete_leader.assert_not_called()
+        assert repository.get_current_leader("wf-1") == "sid-leader"
 
     def test_broadcast_leader_change_logs_emit_errors(
         self,
-        service: tuple[WorkflowCollaborationService, Mock, Mock],
+        real_service: ServiceFixture,
         caplog: pytest.LogCaptureFixture,
     ) -> None:
-        collaboration_service, repository, socketio = service
-        repository.get_session_sids.return_value = ["sid-1", "sid-2"]
+        collaboration_service, repository, socketio = real_service
+        seed_session(repository, "sid-1")
+        seed_session(repository, "sid-2")
         socketio.emit.side_effect = [RuntimeError("boom"), None]
 
         with caplog.at_level(logging.ERROR):
@@ -1149,16 +1149,17 @@ class TestWorkflowCollaborationService:
         assert len(error_records) == 1
         assert "Failed to emit leader status to session sid-1" in error_records[0].getMessage()
 
-    def test_broadcast_online_users_sorts_and_emits(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
+    def test_broadcast_online_users_sorts_and_emits(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, socketio = service
-        repository.list_sessions.return_value = [
-            {"user_id": "u-1", "username": "A", "avatar": None, "sid": "sid-1", "connected_at": 3},
-            {"user_id": "u-2", "username": "B", "avatar": None, "sid": "sid-2", "connected_at": 1},
-        ]
-        repository.get_current_leader.return_value = "sid-1"
+        collaboration_service, repository, socketio = real_service
+        seed_sessions(
+            repository,
+            [
+                {"user_id": "u-1", "username": "A", "avatar": None, "sid": "sid-1", "connected_at": 3},
+                {"user_id": "u-2", "username": "B", "avatar": None, "sid": "sid-2", "connected_at": 1},
+            ],
+        )
+        repository.set_leader("wf-1", "sid-1")
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             # Act
@@ -1178,12 +1179,10 @@ class TestWorkflowCollaborationService:
             room="wf-1",
         )
 
-    def test_broadcast_online_users_reassigns_missing_leader(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, socketio = service
+    def test_broadcast_online_users_reassigns_missing_leader(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, socketio = real_service
         users = [{"user_id": "u-2", "username": "B", "avatar": None, "sid": "sid-2", "connected_at": 1}]
-        repository.get_current_leader.return_value = "sid-old"
+        repository.set_leader("wf-1", "sid-old")
 
         with (
             patch.object(collaboration_service, "_prune_inactive_sessions", return_value=users),
@@ -1192,8 +1191,7 @@ class TestWorkflowCollaborationService:
         ):
             collaboration_service.broadcast_online_users("wf-1")
 
-        repository.delete_leader.assert_called_once_with("wf-1")
-        repository.set_leader.assert_called_once_with("wf-1", "sid-2")
+        assert repository.get_current_leader("wf-1") == "sid-2"
         broadcast_leader_change.assert_called_once_with("wf-1", "sid-2")
         socketio.emit.assert_called_once_with(
             "online_users",
@@ -1202,37 +1200,40 @@ class TestWorkflowCollaborationService:
         )
 
     def test_refresh_session_state_expires_active_leader(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
+        self, real_service: ServiceFixture, collaboration_redis: RedisClientWrapper
     ) -> None:
         # Arrange
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = "sid-1"
+        collaboration_service, repository, _socketio = real_service
+        seed_session(repository, "sid-1")
+        repository.set_leader("wf-1", "sid-1")
+        collaboration_redis.expire(repository.leader_key("wf-1"), 1)
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             # Act
             collaboration_service.refresh_session_state("wf-1", "sid-1")
 
         # Assert
-        repository.refresh_session_state.assert_called_once_with("wf-1", "sid-1")
-        repository.expire_leader.assert_called_once_with("wf-1")
-        repository.set_leader.assert_not_called()
+        assert repository.server_heartbeat_exists("server-1")
+        assert collaboration_redis.ttl(repository.leader_key("wf-1")) == 3600
+        assert repository.get_current_leader("wf-1") == "sid-1"
 
-    def test_refresh_session_state_sets_leader_when_missing(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
+    def test_refresh_session_state_sets_leader_when_missing(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = None
-        repository.list_sessions.return_value = [
-            {
-                "user_id": "u-2",
-                "username": "B",
-                "avatar": None,
-                "sid": "sid-2",
-                "connected_at": 1,
-                "graph_active": True,
-            }
-        ]
+        collaboration_service, repository, _socketio = real_service
+        assert repository.get_current_leader("wf-1") is None
+        seed_sessions(
+            repository,
+            [
+                {
+                    "user_id": "u-2",
+                    "username": "B",
+                    "avatar": None,
+                    "sid": "sid-2",
+                    "connected_at": 1,
+                    "graph_active": True,
+                }
+            ],
+        )
 
         with (
             patch.object(collaboration_service, "is_session_active", return_value=True),
@@ -1242,14 +1243,12 @@ class TestWorkflowCollaborationService:
             collaboration_service.refresh_session_state("wf-1", "sid-2")
 
         # Assert
-        repository.set_leader.assert_called_once_with("wf-1", "sid-2")
+        assert repository.get_current_leader("wf-1") == "sid-2"
         broadcast_leader_change.assert_called_once_with("wf-1", "sid-2")
 
-    def test_refresh_session_state_replaces_inactive_existing_leader(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = "sid-old"
+    def test_refresh_session_state_replaces_inactive_existing_leader(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        repository.set_leader("wf-1", "sid-old")
 
         with (
             patch.object(collaboration_service, "is_session_active", return_value=False),
@@ -1257,33 +1256,31 @@ class TestWorkflowCollaborationService:
         ):
             collaboration_service.refresh_session_state("wf-1", "sid-new")
 
-        repository.delete_leader.assert_called_once_with("wf-1")
-        repository.set_leader.assert_called_once_with("wf-1", "sid-new")
+        assert repository.get_current_leader("wf-1") == "sid-new"
         broadcast_leader_change.assert_called_once_with("wf-1", "sid-new")
 
-    def test_relay_graph_event_emits_update(self, service: tuple[WorkflowCollaborationService, Mock, Mock]) -> None:
+    def test_relay_graph_event_emits_update(self, real_service: ServiceFixture) -> None:
         # Arrange
-        collaboration_service, repository, socketio = service
-        repository.get_sid_mapping.return_value = {"workflow_id": "wf-1", "user_id": "u-1"}
+        collaboration_service, repository, socketio = real_service
+        seed_session(repository, "sid-1")
+        repository.set_leader("wf-1", "sid-1")
 
         # Act
         result = collaboration_service.relay_graph_event("sid-1", {"nodes": []})
 
         # Assert
         assert result == ({"msg": "graph_update_broadcasted"}, 200)
-        repository.refresh_session_state.assert_called_once_with("wf-1", "sid-1")
+        assert repository.server_heartbeat_exists("server-1")
         socketio.emit.assert_called_once_with("graph_update", {"nodes": []}, room="wf-1", skip_sid="sid-1")
 
-    def test_prune_inactive_sessions_handles_empty_and_removes_stale(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.list_sessions.return_value = []
+    def test_prune_inactive_sessions_handles_empty_and_removes_stale(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        seed_sessions(repository, [])
         assert collaboration_service._prune_inactive_sessions("wf-1") == []
 
-        active = {"sid": "sid-1", "user_id": "u-1", "connected_at": 1}
-        stale = {"sid": "sid-2", "user_id": "u-2", "connected_at": 2}
-        repository.list_sessions.return_value = [active, stale]
+        active = {"sid": "sid-1", "user_id": "u-1", "username": "A", "avatar": None, "connected_at": 1}
+        stale = {"sid": "sid-2", "user_id": "u-2", "username": "B", "avatar": None, "connected_at": 2}
+        seed_sessions(repository, [active, stale])
 
         with patch.object(
             collaboration_service,
@@ -1293,53 +1290,67 @@ class TestWorkflowCollaborationService:
             users = collaboration_service._prune_inactive_sessions("wf-1")
 
         assert users == [active]
-        repository.delete_session.assert_called_with("wf-1", "sid-2")
-
-    def test_is_session_active_guard_branches(self, service: tuple[WorkflowCollaborationService, Mock, Mock]) -> None:
-        collaboration_service, repository, socketio = service
-        repository.get_sid_mapping.return_value = {"workflow_id": "wf-1", "user_id": "u-1", "server_id": "server-1"}
-        repository.session_exists.return_value = True
+        assert not repository.session_exists("wf-1", "sid-2")
 
+    def test_is_session_active_guard_branches(
+        self, real_service: ServiceFixture, collaboration_redis: RedisClientWrapper
+    ) -> None:
+        collaboration_service, repository, socketio = real_service
+        info = {
+            "sid": "sid-1",
+            "user_id": "u-1",
+            "username": "A",
+            "avatar": None,
+            "connected_at": 1,
+            "server_id": "server-1",
+        }
+        repository.set_session_info("wf-1", info)
         assert collaboration_service.is_session_active("wf-1", "") is False
-
         socketio.manager.is_connected.return_value = False
         assert collaboration_service.is_session_active("wf-1", "sid-1") is False
-
         socketio.manager.is_connected.return_value = True
         assert collaboration_service.is_session_active("wf-1", "sid-1") is True
-
         socketio.manager.is_connected.side_effect = AttributeError("missing manager")
         assert collaboration_service.is_session_active("wf-1", "sid-1") is False
         socketio.manager.is_connected.side_effect = None
 
-        repository.session_exists.return_value = False
+        collaboration_redis.hdel(repository.workflow_key("wf-1"), "sid-1")
+        assert collaboration_service.is_session_active("wf-1", "sid-1") is False
+        repository.set_session_info("wf-1", info)
+        collaboration_redis.delete(repository.sid_key("sid-1"))
         assert collaboration_service.is_session_active("wf-1", "sid-1") is False
 
-        repository.session_exists.return_value = True
-        repository.get_sid_mapping.return_value = None
-        assert collaboration_service.is_session_active("wf-1", "sid-1") is False
-
-    def test_is_session_active_accepts_remote_session_with_live_server(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, socketio = service
-        repository.get_sid_mapping.return_value = {"workflow_id": "wf-1", "user_id": "u-1", "server_id": "server-2"}
-        repository.session_exists.return_value = True
-        repository.server_heartbeat_exists.return_value = True
+    def test_is_session_active_accepts_remote_session_with_live_server(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, socketio = real_service
+        repository.set_session_info(
+            "wf-1",
+            {
+                "sid": "sid-remote",
+                "user_id": "u-1",
+                "username": "A",
+                "avatar": None,
+                "connected_at": 1,
+                "server_id": "server-2",
+            },
+        )
+        repository.refresh_server_heartbeat("server-2")
         socketio.manager.is_connected.return_value = False
-
         assert collaboration_service.is_session_active("wf-1", "sid-remote") is True
-        repository.server_heartbeat_exists.assert_called_once_with("server-2")
 
-    def test_is_session_active_rejects_remote_session_with_dead_server(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, socketio = service
-        repository.get_sid_mapping.return_value = {"workflow_id": "wf-1", "user_id": "u-1", "server_id": "server-2"}
-        repository.session_exists.return_value = True
-        repository.server_heartbeat_exists.return_value = False
+    def test_is_session_active_rejects_remote_session_with_dead_server(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, socketio = real_service
+        repository.set_session_info(
+            "wf-1",
+            {
+                "sid": "sid-remote",
+                "user_id": "u-1",
+                "username": "A",
+                "avatar": None,
+                "connected_at": 1,
+                "server_id": "server-2",
+            },
+        )
         socketio.manager.is_connected.return_value = False
-
         assert collaboration_service.is_session_active("wf-1", "sid-remote") is False
 
     @staticmethod
@@ -1353,56 +1364,60 @@ class TestWorkflowCollaborationService:
             "graph_active": graph_active,
         }
 
-    def test_select_graph_leader_prefers_visible_preferred(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.list_sessions.return_value = [
-            self._session("sid-1", 1, graph_active=True),
-            self._session("sid-2", 2, graph_active=True),
-        ]
+    def test_select_graph_leader_prefers_visible_preferred(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        seed_sessions(
+            repository,
+            [
+                self._session("sid-1", 1, graph_active=True),
+                self._session("sid-2", 2, graph_active=True),
+            ],
+        )
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             result = collaboration_service._select_graph_leader("wf-1", preferred_sid="sid-2")
 
         assert result == "sid-2"
 
-    def test_select_graph_leader_visible_beats_hidden_preferred(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.list_sessions.return_value = [
-            self._session("sid-1", 1, graph_active=False),
-            self._session("sid-2", 2, graph_active=True),
-        ]
+    def test_select_graph_leader_visible_beats_hidden_preferred(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        seed_sessions(
+            repository,
+            [
+                self._session("sid-1", 1, graph_active=False),
+                self._session("sid-2", 2, graph_active=True),
+            ],
+        )
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             result = collaboration_service._select_graph_leader("wf-1", preferred_sid="sid-1")
 
         assert result == "sid-2"
 
-    def test_select_graph_leader_all_hidden_falls_back_to_preferred(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.list_sessions.return_value = [
-            self._session("sid-1", 1, graph_active=False),
-            self._session("sid-2", 2, graph_active=False),
-        ]
+    def test_select_graph_leader_all_hidden_falls_back_to_preferred(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        seed_sessions(
+            repository,
+            [
+                self._session("sid-1", 1, graph_active=False),
+                self._session("sid-2", 2, graph_active=False),
+            ],
+        )
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             result = collaboration_service._select_graph_leader("wf-1", preferred_sid="sid-2")
 
         assert result == "sid-2"
 
-    def test_select_graph_leader_all_hidden_without_preference_picks_first(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.list_sessions.return_value = [
-            self._session("sid-1", 1, graph_active=False),
-            self._session("sid-2", 2, graph_active=False),
-        ]
+    def test_select_graph_leader_all_hidden_without_preference_picks_first(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        seed_sessions(
+            repository,
+            [
+                self._session("sid-1", 1, graph_active=False),
+                self._session("sid-2", 2, graph_active=False),
+            ],
+        )
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             result = collaboration_service._select_graph_leader("wf-1")
@@ -1410,23 +1425,24 @@ class TestWorkflowCollaborationService:
         assert result == "sid-1"
 
     def test_select_graph_leader_require_graph_active_returns_none_when_all_hidden(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
+        self, real_service: ServiceFixture
     ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.list_sessions.return_value = [
-            self._session("sid-1", 1, graph_active=False),
-        ]
+        collaboration_service, repository, _socketio = real_service
+        seed_sessions(
+            repository,
+            [
+                self._session("sid-1", 1, graph_active=False),
+            ],
+        )
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             result = collaboration_service._select_graph_leader("wf-1", require_graph_active=True)
 
         assert result is None
 
-    def test_select_graph_leader_no_sessions_returns_none(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
-    ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.list_sessions.return_value = []
+    def test_select_graph_leader_no_sessions_returns_none(self, real_service: ServiceFixture) -> None:
+        collaboration_service, repository, _socketio = real_service
+        seed_sessions(repository, [])
 
         with patch.object(collaboration_service, "is_session_active", return_value=True):
             result = collaboration_service._select_graph_leader("wf-1")
@@ -1434,13 +1450,16 @@ class TestWorkflowCollaborationService:
         assert result is None
 
     def test_handle_leader_disconnect_elects_hidden_session_when_no_visible_remains(
-        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
+        self, real_service: ServiceFixture
     ) -> None:
-        collaboration_service, repository, _socketio = service
-        repository.get_current_leader.return_value = "sid-old"
-        repository.list_sessions.return_value = [
-            self._session("sid-2", 2, graph_active=False),
-        ]
+        collaboration_service, repository, _socketio = real_service
+        repository.set_leader("wf-1", "sid-old")
+        seed_sessions(
+            repository,
+            [
+                self._session("sid-2", 2, graph_active=False),
+            ],
+        )
 
         with (
             patch.object(collaboration_service, "is_session_active", return_value=True),
@@ -1448,5 +1467,5 @@ class TestWorkflowCollaborationService:
         ):
             collaboration_service.handle_leader_disconnect("wf-1", "sid-old")
 
-        repository.set_leader.assert_called_once_with("wf-1", "sid-2")
+        assert repository.get_current_leader("wf-1") == "sid-2"
         broadcast_leader_change.assert_called_once_with("wf-1", "sid-2")