Fix unit tests

apache · Aug 22, 2024 · 651d9d5 · 651d9d5
1 parent 6d82ff8
commit 651d9d5
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 12 deletions.
diff --git a/airflow/providers/docker/operators/docker_swarm.py b/airflow/providers/docker/operators/docker_swarm.py
@@ -180,11 +180,14 @@ def _run_service(self) -> None:
                 self.log.info("Service status before exiting: %s", self._service_status())
                 break
 
-        self.tasks = self.cli.tasks(filters={"service": self.service["ID"]})
-        for task in self.tasks:
-            container_id = task["Status"]["ContainerStatus"]["ContainerID"]
-            container = self.cli.inspect_container(container_id)
-            self.containers.append(container)
+        if self.service and self._service_status() == "complete":
+            self.tasks = self.cli.tasks(filters={"service": self.service["ID"]})
+            for task in self.tasks:
+                container_id = task["Status"]["ContainerStatus"]["ContainerID"]
+                container = self.cli.inspect_container(container_id)
+                self.containers.append(container)
+        else:
+            raise AirflowException(f"Service did not complete: {self.service!r}")
 
         if self.retrieve_output:
             return self._attempt_to_retrieve_results()

diff --git a/tests/providers/docker/operators/test_docker_swarm.py b/tests/providers/docker/operators/test_docker_swarm.py
@@ -38,7 +38,7 @@ def _client_tasks_side_effect():
             for _ in range(2):
                 yield [{"Status": {"State": "pending"}}]
             while True:
-                yield [{"Status": {"State": "complete"}}]
+                yield [{"Status": {"State": "complete", "ContainerStatus": {"ContainerID": "some_id"}}}]
 
         def _client_service_logs_effect():
             service_logs = [
@@ -123,7 +123,7 @@ def _client_service_logs_effect():
         assert cskwargs["labels"] == {"name": "airflow__adhoc_airflow__unittest"}
         assert cskwargs["name"].startswith("airflow-")
         assert cskwargs["mode"] == types.ServiceMode(mode="replicated", replicas=3)
-        assert client_mock.tasks.call_count == 6
+        assert client_mock.tasks.call_count == 8
         client_mock.remove_service.assert_called_once_with("some_id")
 
     @mock.patch("airflow.providers.docker.operators.docker_swarm.types")
@@ -134,7 +134,9 @@ def test_auto_remove(self, types_mock, docker_api_client_patcher):
         client_mock.create_service.return_value = {"ID": "some_id"}
         client_mock.images.return_value = []
         client_mock.pull.return_value = [b'{"status":"pull log"}']
-        client_mock.tasks.return_value = [{"Status": {"State": "complete"}}]
+        client_mock.tasks.return_value = [
+            {"Status": {"State": "complete", "ContainerStatus": {"ContainerID": "some_id"}}}
+        ]
         types_mock.TaskTemplate.return_value = mock_obj
         types_mock.ContainerSpec.return_value = mock_obj
         types_mock.RestartPolicy.return_value = mock_obj
@@ -157,7 +159,9 @@ def test_no_auto_remove(self, types_mock, docker_api_client_patcher):
         client_mock.create_service.return_value = {"ID": "some_id"}
         client_mock.images.return_value = []
         client_mock.pull.return_value = [b'{"status":"pull log"}']
-        client_mock.tasks.return_value = [{"Status": {"State": "complete"}}]
+        client_mock.tasks.return_value = [
+            {"Status": {"State": "complete", "ContainerStatus": {"ContainerID": "some_id"}}}
+        ]
         types_mock.TaskTemplate.return_value = mock_obj
         types_mock.ContainerSpec.return_value = mock_obj
         types_mock.RestartPolicy.return_value = mock_obj
@@ -233,7 +237,9 @@ def test_container_resources(self, types_mock, docker_api_client_patcher):
         client_mock.create_service.return_value = {"ID": "some_id"}
         client_mock.images.return_value = []
         client_mock.pull.return_value = [b'{"status":"pull log"}']
-        client_mock.tasks.return_value = [{"Status": {"State": "complete"}}]
+        client_mock.tasks.return_value = [
+            {"Status": {"State": "complete", "ContainerStatus": {"ContainerID": "some_id"}}}
+        ]
         types_mock.TaskTemplate.return_value = mock_obj
         types_mock.ContainerSpec.return_value = mock_obj
         types_mock.RestartPolicy.return_value = mock_obj
@@ -278,7 +284,9 @@ def test_service_args_str(self, types_mock, docker_api_client_patcher):
         client_mock.create_service.return_value = {"ID": "some_id"}
         client_mock.images.return_value = []
         client_mock.pull.return_value = [b'{"status":"pull log"}']
-        client_mock.tasks.return_value = [{"Status": {"State": "complete"}}]
+        client_mock.tasks.return_value = [
+            {"Status": {"State": "complete", "ContainerStatus": {"ContainerID": "some_id"}}}
+        ]
         types_mock.TaskTemplate.return_value = mock_obj
         types_mock.ContainerSpec.return_value = mock_obj
         types_mock.RestartPolicy.return_value = mock_obj
@@ -316,7 +324,9 @@ def test_service_args_list(self, types_mock, docker_api_client_patcher):
         client_mock.create_service.return_value = {"ID": "some_id"}
         client_mock.images.return_value = []
         client_mock.pull.return_value = [b'{"status":"pull log"}']
-        client_mock.tasks.return_value = [{"Status": {"State": "complete"}}]
+        client_mock.tasks.return_value = [
+            {"Status": {"State": "complete", "ContainerStatus": {"ContainerID": "some_id"}}}
+        ]
         types_mock.TaskTemplate.return_value = mock_obj
         types_mock.ContainerSpec.return_value = mock_obj
         types_mock.RestartPolicy.return_value = mock_obj