Merge pull request #88 from GabrielSalla/add-code-module-register-method

Add code module register method

Author: GabrielSalla

src/components/monitors_loader/monitors_loader.py

@@ -94,10 +94,7 @@ async def register_monitor(
 
     monitor = await Monitor.get_or_create(name=monitor_name)
     code_module = await CodeModule.get_or_create(monitor_id=monitor.id)
-    code_module.code = monitor_code
-    code_module.additional_files = additional_files or {}
-    code_module.registered_at = now()
-    await code_module.save()
+    await code_module.register(code=monitor_code, additional_files=additional_files or {})
 
     return monitor
 

src/models/code_module.py

@@ -41,3 +41,11 @@ async def get_updated_code_modules(
             cls.monitor_id.in_(monitors_ids),
             cls.registered_at > reference_timestamp,
         )
+
+    @Base.lock_change
+    async def register(self, code: str, additional_files: dict[str, str] | None = None) -> None:
+        """Register a code module with the given code and additional files"""
+        self.code = code
+        self.additional_files = additional_files or {}
+        self.registered_at = datetime.now()
+        await self.save()

tests/models/test_code_module.py

@@ -1,10 +1,11 @@
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 
 import pytest
 import pytest_asyncio
 
 import databases
-from models import CodeModule
+from models import CodeModule, Monitor
+from utils.time import now
 
 pytestmark = pytest.mark.asyncio(loop_scope="session")
 
@@ -92,3 +93,22 @@ async def test_get_updated_code_modules_timestamp(
     monitors_ids = {code_module.monitor_id for code_module in code_modules}
 
     assert monitors_ids == expected_result
+
+
+async def test_register(sample_monitor: Monitor):
+    """'CodeModule.register' should register a code module with the given code and additional
+    files"""
+    code_module = await CodeModule.get(CodeModule.monitor_id == sample_monitor.id)
+
+    assert code_module is not None
+
+    await code_module.register(
+        code="def get_value(): return 20",
+        additional_files={"file1.py": "content1", "file2.py": "content2"},
+    )
+
+    await code_module.refresh()
+
+    assert code_module.code == "def get_value(): return 20"
+    assert code_module.additional_files == {"file1.py": "content1", "file2.py": "content2"}
+    assert code_module.registered_at > now() - timedelta(seconds=1)