refactor

Author: alemar99

src/maastemporalworker/schedules.py

@@ -12,6 +12,7 @@
     ScheduleOverlapPolicy,
     SchedulePolicy,
     ScheduleSpec,
+    ScheduleState,
     ScheduleUpdate,
     ScheduleUpdateInput,
 )
@@ -43,12 +44,16 @@
             # Will be updated at startup with the value from the db
             intervals=[ScheduleIntervalSpec(every=timedelta(minutes=60))]
         ),
+        # Will be updated at startup with the value from the db
+        state=ScheduleState(paused=True),
         policy=SchedulePolicy(overlap=ScheduleOverlapPolicy.CANCEL_OTHER),
     ),
 }
 
 
-def _master_image_sync_updater(sync_interval_minutes: int):
+def _master_image_sync_updater(
+    sync_interval_minutes: int, auto_import_enabled_config: bool | None = None
+):
     async def do_update(input: ScheduleUpdateInput) -> ScheduleUpdate:
         master_image_sync_schedule = SCHEDULES[MASTER_IMAGE_SYNC_WORKFLOW_NAME]
         master_image_sync_schedule.spec = ScheduleSpec(
@@ -63,6 +68,9 @@ async def do_update(input: ScheduleUpdateInput) -> ScheduleUpdate:
         # Here we save the current state in order to avoid un-pausing the schedule
         # when only changing the sync interval
         current_state = schedule_description.schedule.state
+        if auto_import_enabled_config is not None:
+            # On startup, set the state based on the config
+            current_state.paused = not auto_import_enabled_config
         schedule_description.schedule = master_image_sync_schedule
         schedule_description.schedule.state = current_state
         return ScheduleUpdate(schedule=schedule_description.schedule)
@@ -71,11 +79,15 @@ async def do_update(input: ScheduleUpdateInput) -> ScheduleUpdate:
 
 
 async def update_master_image_sync_schedule_interval(
-    client: Client, sync_interval_minutes_config: int
+    client: Client,
+    sync_interval_minutes_config: int,
+    auto_import_enabled_config: bool | None = None,
 ):
     handle = client.get_schedule_handle(MASTER_IMAGE_SYNC_WORKFLOW_NAME)
     await handle.update(
-        _master_image_sync_updater(sync_interval_minutes_config)
+        _master_image_sync_updater(
+            sync_interval_minutes_config, auto_import_enabled_config
+        )
     )
 
 

src/maastemporalworker/temporal_script.py

@@ -23,7 +23,6 @@
     TemporalServiceCache,
 )
 from maastemporalworker.schedules import (
-    pause_or_unpause_master_image_sync_schedule,
     setup_schedules,
     update_master_image_sync_schedule_interval,
 )
@@ -294,12 +293,12 @@ async def main() -> None:
                 }
             )
 
-    await pause_or_unpause_master_image_sync_schedule(
-        temporal_client, configs[BootImagesAutoImportConfig.name]
-    )
-
     await update_master_image_sync_schedule_interval(
-        temporal_client, configs[BootImagesImportIntervalMinutesConfig.name]
+        temporal_client,
+        sync_interval_minutes_config=configs[
+            BootImagesImportIntervalMinutesConfig.name
+        ],
+        auto_import_enabled_config=configs[BootImagesAutoImportConfig.name],
     )
 
     log.info("temporal-worker started")

src/tests/maastemporalworker/test_schedules.py

@@ -125,12 +125,27 @@ async def test_updates_sync_interval(self):
 
     async def test_preserves_current_state(self):
         mock_state = Mock()
+        mock_state.paused = False
         mock_input = Mock()
         mock_input.description.schedule.state = mock_state
 
         result = await _master_image_sync_updater(30)(mock_input)
 
         assert result.schedule.state is mock_state
+        assert result.schedule.state.paused is False
+
+    async def test_state_changes_based_on_config(self):
+        mock_state = Mock()
+        mock_state.paused = True
+        mock_input = Mock()
+        mock_input.description.schedule.state = mock_state
+
+        result = await _master_image_sync_updater(
+            sync_interval_minutes=30, auto_import_enabled_config=True
+        )(mock_input)
+
+        assert result.schedule.state is mock_state
+        assert result.schedule.state.paused is False
 
 
 @pytest.mark.asyncio