From cf3a36779419f54b5718304df1547de301130142 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 21 Jul 2026 11:12:11 -0300 Subject: [PATCH] SIENTIAPDE-1646 SIENTIAPDE-1646: Refactor TemporalManager by extracting schedule update and task queue name building logic into dedicated helper methods. --- orchestrator/activities/temporal_manager.py | 111 ++++++++++-------- .../activities/test_temporal_manager.py | 85 ++++++++++++++ 2 files changed, 149 insertions(+), 47 deletions(-) diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py index 8768c50..8fcc95c 100644 --- a/orchestrator/activities/temporal_manager.py +++ b/orchestrator/activities/temporal_manager.py @@ -217,12 +217,7 @@ class TemporalManager(SientiaMonitoring): f'{json.dumps(schedule, indent=4, sort_keys=True)}', metadata=metadata ) - runtime_name = ( - schedule.get('runtime', 'legacy') - if workflow_type in RUNTIME_WORKFLOWS - else None - ) - task_queue_name = build_queue_name(workflow_type, runtime_name) + task_queue_name = self._build_task_queue_name(workflow_type, schedule) schedule['task_queue'] = task_queue_name await client.create_schedule( @@ -285,6 +280,68 @@ class TemporalManager(SientiaMonitoring): return report + @staticmethod + def _build_task_queue_name(workflow_type: str, schedule: dict[str, Any]) -> str: + """ + Build the Temporal task queue name for a schedule. + + Only workflow types listed in ``RUNTIME_WORKFLOWS`` get an environment/tenant + specific ``runtime`` suffix; every other workflow type gets a plain queue name. + """ + runtime_name = ( + schedule.get('runtime', 'legacy') if workflow_type in RUNTIME_WORKFLOWS else None + ) + return build_queue_name(workflow_type, runtime_name) + + def _make_schedule_updater(self, schedule: dict[str, Any], metadata: dict[str, Any]): + """Build the ``ScheduleUpdate`` callback used by ``handler.update`` for one schedule.""" + + # fmt: off + async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR + schedule_action = input_data.description.schedule.action + + self.debug("Updating schedule:", metadata=metadata) + + if hasattr(schedule_action, "args"): + self.debug("New schedule:", metadata=metadata) + self.debug( + f"{json.dumps(schedule, indent=4, sort_keys=True)}", metadata=metadata) # NOSONAR + + schedule_action.args = [schedule] + + input_data.description.schedule.spec.intervals = [ + ScheduleIntervalSpec( + every=timedelta( + seconds=parse_frequency(schedule.get('frequency', '1m'))), + offset=timedelta( + seconds=parse_frequency(schedule.get('offset', '0m'))), + ) + ] + + return ScheduleUpdate(schedule=input_data.description.schedule) + + # fmt: on + return update_schedule + + async def _update_single_schedule( + self, + client: Client, + schedule_name: str, + schedule: dict[str, Any], + metadata: dict[str, Any], + ) -> None: + """Update a single schedule in Temporal, raising if the schedule handle is missing.""" + handler = client.get_schedule_handle(schedule_name) + + if not handler: + raise ValueError(f'Schedule {schedule_name} not found') + + workflow_type = schedule['workflow_type'] + schedule['task_queue'] = self._build_task_queue_name(workflow_type, schedule) + + update_schedule = self._make_schedule_updater(schedule, metadata) + await handler.update(update_schedule) + @activity.defn(name='update_schedules') async def update_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: """ @@ -321,47 +378,7 @@ class TemporalManager(SientiaMonitoring): for schedule_name, schedule in schedules.items(): try: - handler = client.get_schedule_handle(schedule_name) - - if not handler: - raise ValueError(f'Schedule {schedule_name} not found') - - workflow_type = schedule['workflow_type'] - runtime_name = ( - schedule.get('runtime', 'legacy') - if workflow_type in RUNTIME_WORKFLOWS - else None - ) - schedule['task_queue'] = build_queue_name(workflow_type, runtime_name) - - # fmt: off - async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR - schedule_action = input_data.description.schedule.action - - self.debug("Updating schedule:", metadata=metadata) - - if hasattr(schedule_action, "args"): - self.debug("New schedule:", metadata=metadata) - self.debug( - f"{json.dumps(schedule, indent=4, sort_keys=True)}", metadata=metadata) # NOSONAR - - schedule_action.args = [schedule] - - input_data.description.schedule.spec.intervals = [ - ScheduleIntervalSpec( - every=timedelta( - seconds=parse_frequency(schedule.get('frequency', '1m'))), - offset=timedelta( - seconds=parse_frequency(schedule.get('offset', '0m'))), - ) - ] - - return ScheduleUpdate(schedule=input_data.description.schedule) - - # fmt: on - await handler.update(update_schedule) - - del update_schedule + await self._update_single_schedule(client, schedule_name, schedule, metadata) report.append( { diff --git a/tests/orchestrator/activities/test_temporal_manager.py b/tests/orchestrator/activities/test_temporal_manager.py index 07932ce..337ff1a 100644 --- a/tests/orchestrator/activities/test_temporal_manager.py +++ b/tests/orchestrator/activities/test_temporal_manager.py @@ -504,6 +504,91 @@ async def test_delete_schedules_with_no_client(temporal_manager): ) +def test_build_task_queue_name_non_runtime_workflow(temporal_manager): + task_queue_name = temporal_manager._build_task_queue_name( + 'test-workflow', {'workflow_type': 'test-workflow'} + ) + + assert task_queue_name == 'test-workflow-queue' + + +def test_build_task_queue_name_default_runtime_legacy(temporal_manager): + task_queue_name = temporal_manager._build_task_queue_name('drift', {'workflow_type': 'drift'}) + + assert task_queue_name == 'drift-legacy-queue' + + +def test_build_task_queue_name_explicit_runtime(temporal_manager): + task_queue_name = temporal_manager._build_task_queue_name( + 'drift', {'workflow_type': 'drift', 'runtime': 'tenant-x'} + ) + + assert task_queue_name == 'drift-tenant-x-queue' + + +@mark.asyncio +@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency) +@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec') +async def test_make_schedule_updater_patches_args_and_intervals( + _mock_schedule_interval_spec, _mock_parse_frequency, temporal_manager +): + schedule = {'workflow_type': 'scouter', 'frequency': '2m', 'offset': '1h', 'data': 'new'} + update_schedule = temporal_manager._make_schedule_updater(schedule, metadata) + + schedule_action = MagicMock() + input_data = MagicMock() + input_data.description.schedule.action = schedule_action + + result = await update_schedule(input_data) + + assert schedule_action.args == [schedule] + assert input_data.description.schedule.spec.intervals == [ + _mock_schedule_interval_spec.return_value + ] + _mock_schedule_interval_spec.assert_called_once_with( + every=timedelta(seconds=120), offset=timedelta(seconds=3600) + ) + assert result.schedule == input_data.description.schedule + + +@mark.asyncio +async def test_make_schedule_updater_skips_args_without_attribute(temporal_manager): + schedule = {'workflow_type': 'scouter', 'frequency': '1m', 'data': 'new'} + update_schedule = temporal_manager._make_schedule_updater(schedule, metadata) + + input_data = MagicMock() + input_data.description.schedule.action = object() + + await update_schedule(input_data) + + +@mark.asyncio +async def test_update_single_schedule_not_found(temporal_manager): + client = MagicMock(get_schedule_handle=MagicMock(return_value=None)) + + try: + await temporal_manager._update_single_schedule( + client, 'missing-schedule', {'workflow_type': 'scouter'}, metadata + ) + except ValueError as e: + assert str(e) == 'Schedule missing-schedule not found' + else: + raise AssertionError('Expected a ValueError to be raised') + + +@mark.asyncio +async def test_update_single_schedule_updates_handle(temporal_manager): + handler = MagicMock(update=AsyncMock()) + client = MagicMock(get_schedule_handle=MagicMock(return_value=handler)) + schedule = {'workflow_type': 'drift', 'frequency': '1m'} + + await temporal_manager._update_single_schedule(client, 'test-schedule', schedule, metadata) + + client.get_schedule_handle.assert_called_once_with('test-schedule') + handler.update.assert_awaited_once() + assert schedule['task_queue'] == 'drift-legacy-queue' + + @mark.asyncio @patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency) @patch('orchestrator.activities.temporal_manager.Schedule')