monitors loader uses task manager

Author: GabrielSalla

src/components/monitors_loader/__init__.py

@@ -3,13 +3,13 @@
     check_monitor,
     init,
     register_monitor,
-    wait_stop,
+    run,
 )
 
 __all__ = [
     "MonitorValidationError",
     "check_monitor",
     "init",
     "register_monitor",
-    "wait_stop",
+    "run",
 ]

src/components/monitors_loader/monitors_loader.py

@@ -3,7 +3,7 @@
 import shutil
 from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Any, Generator, cast
+from typing import Generator, cast
 
 from pydantic.dataclasses import dataclass
 
@@ -27,8 +27,6 @@
 EARLY_LOAD_TIME = 5
 COOL_DOWN_TIME = 2
 
-_task: asyncio.Task[Any]
-
 
 @dataclass
 class AdditionalFile:
@@ -289,7 +287,7 @@ async def _load_monitors(last_load_time: datetime | None) -> None:
     registry.monitors_pending.clear()
 
 
-async def _run() -> None:
+async def run() -> None:
     """Monitors loading loop, loading them recurrently. Stops automatically when the app stops"""
     last_load_time: datetime | None = None
 
@@ -326,21 +324,13 @@ async def _run() -> None:
             if time_since_last_load < COOL_DOWN_TIME:
                 await app.sleep(COOL_DOWN_TIME - time_since_last_load)
 
+    _logger.info("Removing temporary monitors paths")
+    shutil.rmtree(Path(RELATIVE_PATH) / MONITORS_LOAD_PATH, ignore_errors=True)
+    shutil.rmtree(Path(RELATIVE_PATH) / MONITORS_PATH, ignore_errors=True)
+
 
 async def init(controller_enabled: bool) -> None:
     """Load the internal monitors and sample monitors if controller is enabled, and start the
     monitors load task"""
     if controller_enabled:
         await _register_monitors()
-
-    global _task
-    _task = asyncio.create_task(_run())
-
-
-async def wait_stop() -> None:
-    """Wait for the Monitors load task to finish"""
-    global _task
-    await _task
-    _logger.info("Removing temporary monitors paths")
-    shutil.rmtree(Path(RELATIVE_PATH) / MONITORS_LOAD_PATH, ignore_errors=True)
-    shutil.rmtree(Path(RELATIVE_PATH) / MONITORS_PATH, ignore_errors=True)

src/main.py

@@ -44,7 +44,6 @@ async def init(controller_enabled: bool, executor_enabled: bool) -> None:
 async def finish(controller_enabled: bool, executor_enabled: bool) -> None:
     """Finish the application, making sure any exception won't impact other closing tasks"""
     await protected_task(_logger, http_server.wait_stop())
-    await protected_task(_logger, monitors_loader.wait_stop())
     await protected_task(_logger, databases.close())
     await protected_task(_logger, internal_database.close())
     await protected_task(

tests/components/monitors_loader/test_monitors_loader.py

@@ -820,18 +820,18 @@ async def test_run_as_controller(mocker, monkeypatch, clear_database):
     _load_monitors_spy: AsyncMock = mocker.spy(monitors_loader, "_load_monitors")
 
     await monitors_loader.init(controller_enabled=True)
+    run_task = asyncio.create_task(monitors_loader.run())
 
     for _ in range(3):
         await asyncio.sleep(0.1)
         registry.monitors_pending.set()
 
     await asyncio.sleep(0.1)
     app.stop()
-    await monitors_loader.wait_stop()
 
     assert _load_monitors_spy.call_count == 4
 
-    assert monitors_loader._task.done()
+    await asyncio.wait_for(run_task, timeout=0.1)
 
     assert len(registry._monitors) == 2
 
@@ -862,18 +862,18 @@ async def test_run_as_executor(mocker, monkeypatch, clear_database):
     )
 
     await monitors_loader.init(controller_enabled=False)
+    run_task = asyncio.create_task(monitors_loader.run())
 
     for _ in range(3):
         await asyncio.sleep(0.1)
         registry.monitors_pending.set()
 
     await asyncio.sleep(0.1)
     app.stop()
-    await monitors_loader.wait_stop()
 
     assert _load_monitors_spy.call_count == 4
 
-    assert monitors_loader._task.done()
+    await asyncio.wait_for(run_task, timeout=0.1)
 
     assert len(registry._monitors) == 2
     assert isinstance(registry._monitors[9999123]["module"], ModuleType)
@@ -905,18 +905,18 @@ async def test_run_cool_down(mocker, monkeypatch, clear_database):
     )
 
     await monitors_loader.init(controller_enabled=False)
+    run_task = asyncio.create_task(monitors_loader.run())
 
     for _ in range(3):
         await asyncio.sleep(0.1)
         registry.monitors_pending.set()
 
     await asyncio.sleep(0.1)
     app.stop()
-    await monitors_loader.wait_stop()
 
     assert _load_monitors_spy.call_count == 1
 
-    assert monitors_loader._task.done()
+    await asyncio.wait_for(run_task, timeout=0.1)
 
     assert len(registry._monitors) == 2
     assert isinstance(registry._monitors[9999123]["module"], ModuleType)
@@ -959,10 +959,10 @@ async def test_run_sleep_time(mocker, monkeypatch, clear_database, seconds, expe
     )
 
     await monitors_loader.init(controller_enabled=False)
+    run_task = asyncio.create_task(monitors_loader.run())
     await asyncio.sleep(1)
     app.stop()
-    await monitors_loader.wait_stop()
 
-    assert monitors_loader._task.done()
+    await asyncio.wait_for(run_task, timeout=0.1)
 
     sleep_spy.assert_called_once_with(expected_sleep_time)