Merge pull request #789 from Ostorlab/fix/exit-once-the-scan-has-comp…

…leted

Fix/exit once the scan has completed

src/ostorlab/runtimes/local/log_streamer.py

@@ -40,6 +40,7 @@ class LogStream:
     def __init__(self):
         self._threads = []
         self._color_map = {}
+        self.services = []
 
     def stream(self, service: docker.models.services.Service) -> None:
         """Stream logs of a service without blocking.
@@ -51,8 +52,9 @@ def stream(self, service: docker.models.services.Service) -> None:
         """
         color = self._select_color(service)
         logs = service.logs(details=False, follow=True, stdout=True, stderr=True)
+        self.services.append(service.id)
         t = threading.Thread(
-            target=_stream_log, args=(logs, service.name, color), daemon=False
+            target=_stream_log, args=(logs, service.name, color), daemon=True
         )
         self._threads.append(t)
         t.start()

src/ostorlab/runtimes/local/runtime.py

@@ -5,6 +5,8 @@
 """
 
 import logging
+import sys
+import threading
 from concurrent import futures
 from typing import Dict, List
 from typing import Optional
@@ -232,6 +234,10 @@ def scan(
             console.info("Updating scan status")
             self._update_scan_progress("IN_PROGRESS")
             console.success("Scan created successfully")
+            scan_complete_thread = threading.Thread(
+                target=self._check_services_running, daemon=False
+            )
+            scan_complete_thread.start()
             return self._scan_db
         except AgentNotHealthy:
             message = "Agent not starting"
@@ -254,6 +260,23 @@ def scan(
             self.stop(str(self._scan_db.id))
             raise MissingAgentDefinition(message)
 
+    def _check_services_running(self) -> None:
+        """Check if the services are still running."""
+        if len(self.follow) == 0:
+            return
+
+        stop_event = threading.Event()
+        while stop_event.is_set() is False:
+            for service_id in list(self._log_streamer.services):
+                try:
+                    self._docker_client.services.get(service_id)
+                except docker_errors.NotFound:
+                    self._log_streamer.services.remove(service_id)
+            if len(self._log_streamer.services) == 0:
+                console.success("Scan done.")
+                stop_event.set()
+        sys.exit(0)
+
     def stop(self, scan_id: str) -> None:
         """Remove a service (scan) belonging to universe with scan_id(Universe Id).
 

src/ostorlab/runtimes/runtime.py

@@ -23,7 +23,7 @@ class Scan:
 class Runtime(abc.ABC):
     """Runtime is in charge of preparing the environment to trigger a scan."""
 
-    follow: bool
+    follow: list
 
     @abc.abstractmethod
     def can_run(self, agent_group_definition: definitions.AgentGroupDefinition) -> bool:

tests/runtimes/local/runtime_test.py

@@ -1,5 +1,6 @@
 """Unittest for local runtime."""
 
+from time import sleep
 from typing import Any
 
 import docker
@@ -332,3 +333,34 @@ def testRuntimeScanList_whenDockerIsDown_DoesNotCrash(
     scans = local_runtime.LocalRuntime().list()
 
     assert scans is None
+
+
+def testCheckServicesMethod_whenServicesAreStopped_shouldExit(
+    mocker: plugin.MockerFixture, run_scan_mock: None
+) -> None:
+    """Unit test for Check services method when services are stopped"""
+    mocker.patch(
+        "ostorlab.runtimes.definitions.AgentSettings.container_image",
+        return_value="agent_42_docker_image",
+        new_callable=mocker.PropertyMock,
+    )
+    mocker.patch("ostorlab.runtimes.local.agent_runtime.AgentRuntime")
+    mocker.patch("ostorlab.runtimes.local.services.mq.LocalRabbitMQ.is_service_healthy")
+    mocker.patch("ostorlab.runtimes.local.services.redis.LocalRedis.is_service_healthy")
+    exit_mock = mocker.patch("sys.exit")
+    local_runtime_instance = local_runtime.LocalRuntime()
+    local_runtime_instance.follow = ["service1"]
+    agent_group_definition = definitions.AgentGroupDefinition(
+        agents=[definitions.AgentSettings(key="agent/ostorlab/agent42")]
+    )
+
+    local_runtime_instance.scan(
+        title="test local",
+        agent_group_definition=agent_group_definition,
+        assets=[android_apk.AndroidApk(content=b"APK")],
+    )
+    sleep(1)
+
+    assert exit_mock.call_count == 1
+    exit_with = exit_mock.call_args_list[0][0][0]
+    assert exit_with == 0