Add comprehensive cleanup for local worker upon successful job run and upload

cli/main.py

@@ -42,8 +42,11 @@ def main() -> None:
     worker_info = info.sys.collect()
 
     while True:
+        container = None
+        job_id = None
+        path_job = None
+
         try:
-            container = None
             jobs = api_worker.fetch_jobs(
                 limit=1,
                 cpu_cores=worker_info.sys.cores,
@@ -177,8 +180,41 @@ def main() -> None:
             else:
                 raise e
 
-        logger.info(f"Job {job_id} finished")
-        shutil.rmtree(path_job)
+        finally:
+            # Clean up resources after successful job run AND upload
+            if job_id:
+                logger.info(f"Cleaning up job {job_id}")
+
+                # Clean up Docker container
+                if container:
+                    try:
+                        if container.status == "running":
+                            logger.info(f"Stopping running container for job {job_id}")
+                            container.stop()
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to stop container for job {job_id}: {e}"
+                        )
+
+                    try:
+                        logger.info(f"Removing container for job {job_id}")
+                        container.remove()
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to remove container for job {job_id}: {e}"
+                        )
+
+                # Clean up job directory
+                if path_job and path_job.exists():
+                    try:
+                        logger.info(f"Removing job directory {path_job}")
+                        shutil.rmtree(path_job)
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to clean up job directory {path_job}: {e}"
+                        )
+
+                logger.info(f"Job {job_id} cleanup completed")
 
 
 if __name__ == "__main__":

fetcher/api/__init__.pyi

@@ -1 +1,4 @@
-# ...existing code...
+from . import builder as builder
+from . import model as model
+from . import token as token
+from . import worker as worker

fetcher/docker/__init__.pyi

@@ -1 +0,0 @@
-# ...existing code...

fetcher/info/__init__.pyi

@@ -1 +1 @@
-# ...existing code...
+from . import sys as sys

fetcher/io/__init__.pyi

@@ -1 +1 @@
-# ...existing code...
+from . import files as files

fetcher/status/__init__.pyi

@@ -1 +1,2 @@
-# ...existing code...
+from . import pinger as pinger
+from . import status as status

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fetcher"
-version = "0.3.0"
+version = "0.4.0"
 description = "Job fetcher for workers of DECODE OpenCloud."
 authors = ["Arthur Jaques <arthur.jaques@hispeed.ch>"]
 readme = "README.md"

tests/unit/test_cleanup.py

@@ -0,0 +1,86 @@
+"""Tests for cleanup functionality in the main job processing loop."""
+
+import shutil
+import tempfile
+from pathlib import Path
+from unittest.mock import Mock
+
+
+class TestJobCleanup:
+    """Test cleanup behavior for local worker jobs."""
+
+    def test_container_cleanup_logic(self) -> None:
+        """Test the container cleanup logic works correctly."""
+        # Mock a container
+        mock_container = Mock()
+        mock_container.status = "running"
+
+        job_id = "test_job_123"
+
+        # Test cleanup logic (simulating the finally block)
+        if job_id:
+            if mock_container:
+                try:
+                    if mock_container.status == "running":
+                        mock_container.stop()
+                    mock_container.remove()
+                except Exception:
+                    # Should handle exceptions gracefully
+                    pass
+
+        # Verify methods were called
+        mock_container.stop.assert_called_once()
+        mock_container.remove.assert_called_once()
+
+    def test_container_cleanup_handles_exceptions(self) -> None:
+        """Test that container cleanup handles exceptions gracefully."""
+        # Mock a container that raises exceptions
+        mock_container = Mock()
+        mock_container.status = "running"
+        mock_container.stop.side_effect = Exception("Stop failed")
+        mock_container.remove.side_effect = Exception("Remove failed")
+
+        job_id = "test_job_456"
+
+        # Test cleanup logic should handle each step independently
+        if job_id:
+            if mock_container:
+                try:
+                    if mock_container.status == "running":
+                        mock_container.stop()
+                except Exception:
+                    pass
+                try:
+                    mock_container.remove()
+                except Exception:
+                    pass
+
+        # Verify methods were called despite exceptions
+        mock_container.stop.assert_called_once()
+        mock_container.remove.assert_called_once()
+
+    def test_cleanup_defensive_programming(self) -> None:
+        """Test that cleanup handles missing variables gracefully."""
+        # This test verifies that the cleanup code can handle cases where
+        # container, job_id, or path_job might be None or undefined
+
+        # Create a temporary directory for testing
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path_job = Path(temp_dir) / "test_job"
+            path_job.mkdir()
+
+            # Test cleanup with missing container (should not crash)
+            container = None
+            job_id = "test_job"
+
+            # Simulate the cleanup code
+            if job_id:
+                if container:
+                    # This should not execute
+                    assert False, "Should not try to clean up None container"
+
+                if path_job and path_job.exists():
+                    shutil.rmtree(path_job)
+
+            # Verify directory was cleaned up
+            assert not path_job.exists()