get skill template container

skill-manager/skill_manager/core/docker_client.py

@@ -2,7 +2,7 @@
 import shutil
 from contextlib import contextmanager
 from time import time
-from typing import List
+from typing import List, Union
 
 from docker import DockerClient
 from docker.models.containers import Container
@@ -103,10 +103,16 @@ def _wait_for_container_status(
 
     def get_skill_template_containers(self) -> List[Container]:
         return self.docker_client.containers.list(
-            filters={"labels": "type=skill-template"}
+            filters={"label": "type=skill-template"}
         )
 
-    def get_skill_template_container_by_id(self, skill_template_id) -> Container:
-        return self.docker_client.containers.list(
-            filters={"labels": f"skill-template-id={skill_template_id}"}
-        )[0]
+    def get_skill_template_container_by_id(self, skill_template_id) -> Union[Container, None]:
+        skill_template_container = self.docker_client.containers.list(
+            filters={"label": f"skill-template-id={skill_template_id}"}
+        )
+        if skill_template_container:
+            skill_template_container = skill_template_container[0]
+        else:
+            skill_template_container = None
+
+        return skill_template_container

skill-manager/tests/test_docker_client.py

@@ -219,3 +219,63 @@ async def predict(query):
     # clean-up
     container.stop()
     container.remove()
+
+
+@pytest.mark.parametrize("num_containers", [0, 3], ids=["zero-runnning", "n-running"])
+def test_get_skill_template_containers(num_containers, docker_client: DockerClient):
+
+    test_image = "containous/whoami"
+    containers = []
+    for _ in range(num_containers):
+        containers.append(
+            docker_client.containers.run(
+                test_image, detach=True, remove=True, labels={"type": "skill-template"}
+            )
+        )
+    # start some "distractor" containers
+    for _ in range(4):
+        containers.append(
+            docker_client.containers.run(test_image, detach=True, remove=True)
+        )
+
+    sm_docker_client = SkillManagerDockerClient()
+    found_containers = sm_docker_client.get_skill_template_containers()
+    assert len(found_containers) == num_containers
+
+    # clean-up
+    for c in containers:
+        c.stop()
+
+
+@pytest.mark.parametrize(
+    "skill_template_running", [False, True], ids=["not-runnning", "running"]
+)
+def test_get_skill_template_container_by_id(
+    skill_template_running, docker_client: DockerClient
+):
+
+    skill_template_id = str(uuid.uuid1())
+    test_image = "containous/whoami"
+    containers = []
+    if skill_template_running:
+        containers.append(
+            docker_client.containers.run(
+                test_image, detach=True, remove=True, labels={"skill-template-id": skill_template_id}
+            )
+        )
+    # start some "distractor" containers
+    for _ in range(4):
+        containers.append(
+            docker_client.containers.run(test_image, detach=True, remove=True)
+        )
+
+    sm_docker_client = SkillManagerDockerClient()
+    found_container = sm_docker_client.get_skill_template_container_by_id(skill_template_id)
+    if skill_template_running:
+        assert found_container.labels["skill-template-id"] == skill_template_id
+    else:
+        assert found_container is None
+
+    # clean-up
+    for c in containers:
+        c.stop()