Added unit test setup for worker

.github/workflows/ci-test.yaml

@@ -0,0 +1,57 @@
+name: Run tox tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+    branches: [main]
+
+jobs:
+  test:
+    if: github.event.pull_request.draft == false
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v3
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.9'
+
+    - name: Cache tox environments
+      uses: actions/cache@v3
+      with:
+        path: .tox/
+        key: ${{ runner.os }}-tox-${{ hashFiles('**/pyproject.toml', '**/tox.ini') }}
+        restore-keys: |
+          ${{ runner.os }}-tox-
+
+    - name: Install tox
+      run: pip install tox
+
+    - name: Run tox
+      id: tox
+      run: |
+        tox
+
+    - name: Render the report to the PR
+      uses: marocchino/sticky-pull-request-comment@v2
+      with:
+        header: worker-test-report
+        recreate: true
+        path: worker-report.md
+
+    - name: Output reports to the job summary when tests fail
+      shell: bash
+      run: |
+        if [ -f "worker-report.md" ]; then
+          echo "<details><summary>Worker Test Report</summary>" >> $GITHUB_STEP_SUMMARY
+          echo "" >> $GITHUB_STEP_SUMMARY
+          cat "worker-report.md" >> $GITHUB_STEP_SUMMARY
+          echo "" >> $GITHUB_STEP_SUMMARY
+          echo "</details>" >> $GITHUB_STEP_SUMMARY
+        fi

.gitignore

@@ -53,6 +53,7 @@ pip-delete-this-directory.txt
 
 # Unit test / coverage reports
 htmlcov/
+**/report.html
 .tox/
 .nox/
 .coverage

tox.ini

@@ -0,0 +1,18 @@
+[tox]
+env_list = py{39,310,311}, worker
+
+# [testenv]
+# skip_install = true
+
+[testenv:worker]
+changedir = worker
+setenv =
+    PDM_IGNORE_SAVED_PYTHON="1"
+deps = pdm
+allowlist_externals=
+    sh
+commands_pre =
+    pdm sync --dev
+    sh -c '[ -f cloud_requirements.txt ] && pip install -r cloud_requirements.txt || echo "cloud_requirements.txt not found"'
+commands =
+    pytest -v --md-report-verbose=1 --md-report --md-report-flavor gfm --md-report-output ../worker-report.md

worker/.gitignore

@@ -166,3 +166,4 @@ src/unstract/worker/clients/
 !src/unstract/worker/clients/interface.py
 !src/unstract/worker/clients/helper.py
 !src/unstract/worker/clients/docker.py
+!src/unstract/worker/clients/test_docker.py

worker/pyproject.toml

@@ -29,6 +29,13 @@ deploy = [
     "gunicorn>=21.2.0",
 ]
 
+dev = [
+    "pytest>=8.2.2",
+    "pytest-mock>=3.14.0",
+    "pytest-cov>=5.0.0",
+    "pytest-md-report>=0.6.2",
+]
+
 [tool.pdm.scripts]
 worker.cmd = "flask --app src/unstract/worker/main.py run --port 5002"
 worker.env_file = ".env"

worker/src/unstract/worker/clients/docker.py

@@ -5,13 +5,15 @@
 
 from docker.errors import APIError, ImageNotFound
 from docker.models.containers import Container
+from unstract.worker.clients.helper import ContainerClientHelper
+from unstract.worker.clients.interface import (
+    ContainerClientInterface,
+    ContainerInterface,
+)
 from unstract.worker.constants import Env
 from unstract.worker.utils import Utils
 
-from docker import DockerClient, from_env
-
-from .helper import normalize_container_name
-from .interface import ContainerClientInterface, ContainerInterface
+from docker import DockerClient
 
 
 class DockerContainer(ContainerInterface):
@@ -44,7 +46,7 @@ def __init__(self, image_name: str, image_tag: str, logger: logging.Logger) -> N
 
         # Create a Docker client that communicates with
         #   the Docker daemon in the host environment
-        self.client: DockerClient = from_env()
+        self.client: DockerClient = DockerClient.from_env()
         self.__private_login()
 
     def __private_login(self):
@@ -179,7 +181,7 @@ def get_container_run_config(
                 }
             )
         return {
-            "name": normalize_container_name(self.image_name),
+            "name": ContainerClientHelper.normalize_container_name(self.image_name),
             "image": self.get_image(),
             "command": command,
             "detach": True,

worker/src/unstract/worker/clients/helper.py

@@ -5,11 +5,15 @@
 from .interface import ContainerClientInterface
 
 
-def get_container_client() -> ContainerClientInterface:
-    client_path = os.getenv("CONTAINER_CLIENT_PATH", "unstract.worker.clients.docker")
-    print("Loading the container client from path:", client_path)
-    return import_module(client_path).Client
+class ContainerClientHelper:
+    @staticmethod
+    def get_container_client() -> ContainerClientInterface:
+        client_path = os.getenv(
+            "CONTAINER_CLIENT_PATH", "unstract.worker.clients.docker"
+        )
+        print("Loading the container client from path:", client_path)
+        return import_module(client_path).Client
 
-
-def normalize_container_name(name: str) -> str:
-    return name.replace("/", "-") + "-" + str(uuid.uuid4())
+    @staticmethod
+    def normalize_container_name(name: str) -> str:
+        return name.replace("/", "-") + "-" + str(uuid.uuid4())

worker/src/unstract/worker/clients/test_docker.py

@@ -0,0 +1,195 @@
+import logging
+import os
+from unittest.mock import MagicMock
+
+import pytest
+from docker.errors import ImageNotFound
+from unstract.worker.constants import Env
+
+from .docker import Client, DockerContainer
+
+DOCKER_MODULE = "unstract.worker.clients.docker"
+
+
+@pytest.fixture
+def docker_container():
+    container = MagicMock()
+    return DockerContainer(container)
+
+
+@pytest.fixture
+def docker_client():
+    image_name = "test-image"
+    image_tag = "latest"
+    logger = logging.getLogger("test-logger")
+    return Client(image_name, image_tag, logger)
+
+
+def test_logs(docker_container, mocker):
+    """Test the logs method to ensure it yields log lines."""
+    mock_container = mocker.patch.object(docker_container, "container")
+    mock_container.logs.return_value = [b"log line 1", b"log line 2"]
+
+    logs = list(docker_container.logs(follow=True))
+    assert logs == ["log line 1", "log line 2"]
+
+
+def test_cleanup(docker_container, mocker):
+    """Test the cleanup method to ensure it removes the container."""
+    mock_container = mocker.patch.object(docker_container, "container")
+    mocker.patch(f"{DOCKER_MODULE}.Utils.remove_container_on_exit", return_value=True)
+
+    docker_container.cleanup()
+    mock_container.remove.assert_called_once_with(force=True)
+
+
+def test_cleanup_skip(docker_container, mocker):
+    """Test the cleanup method to ensure it doesn't remove the container."""
+    mock_container = mocker.patch.object(docker_container, "container")
+    mocker.patch(f"{DOCKER_MODULE}.Utils.remove_container_on_exit", return_value=False)
+
+    docker_container.cleanup()
+    mock_container.remove.assert_not_called()
+
+
+def test_client_init(mocker):
+    """Test the Client initialization."""
+    mock_from_env = mocker.patch(f"{DOCKER_MODULE}.DockerClient.from_env")
+    client_instance = Client("test-image", "latest", logging.getLogger("test-logger"))
+
+    mock_from_env.assert_called_once()
+    assert client_instance.client is not None
+
+
+def test_get_image_exists(docker_client, mocker):
+    """Test the __image_exists method."""
+    # Mock the client object
+    mock_client = mocker.patch.object(docker_client, "client")
+    # Create a mock for the 'images' attribute
+    mock_images = mocker.MagicMock()
+    # Attach the mock to the client object
+    mock_client.images = mock_images
+    # Patch the 'get' method of the 'images' attribute
+    mock_images.get.side_effect = ImageNotFound("Image not found")
+
+    assert not docker_client._Client__image_exists("test-image:latest")
+    mock_images.get.assert_called_once_with("test-image:latest")
+
+
+def test_get_image(docker_client, mocker):
+    """Test the get_image method."""
+    # Patch the client object to control its behavior
+    mock_client = mocker.patch.object(docker_client, "client")
+    # Patch the images attribute of the client to control its behavior
+    mock_images = mocker.MagicMock()
+    mock_client.images = mock_images
+
+    # Case 1: Image exists
+    mock_images.get.side_effect = MagicMock()  # Mock that image exists
+    assert docker_client.get_image() == "test-image:latest"
+    mock_images.get.assert_called_once_with("test-image:latest")  # Ensure get is called
+
+    # Case 2: Image does not exist
+    mock_images.get.side_effect = ImageNotFound(
+        "Image not found"
+    )  # Mock that image doesn't exist
+    mock_pull = mocker.patch.object(
+        docker_client.client.api, "pull"
+    )  # Patch pull method
+    mock_pull.return_value = iter([{"status": "pulling"}])  # Simulate pull process
+    assert docker_client.get_image() == "test-image:latest"
+    mock_pull.assert_called_once_with(
+        repository="test-image",
+        tag="latest",
+        stream=True,
+        decode=True,
+    )
+
+
+def test_get_container_run_config(docker_client, mocker):
+    """Test the get_container_run_config method."""
+    os.environ[Env.WORKFLOW_DATA_DIR] = "/source"
+    command = ["echo", "hello"]
+    organization_id = "org123"
+    workflow_id = "wf123"
+    execution_id = "ex123"
+
+    mocker.patch.object(docker_client, "_Client__image_exists", return_value=True)
+    mocker_normalize = mocker.patch(
+        f"{DOCKER_MODULE}.ContainerClientHelper.normalize_container_name",
+        return_value="test-image",
+    )
+    config = docker_client.get_container_run_config(
+        command,
+        organization_id,
+        workflow_id,
+        execution_id,
+        envs={"KEY": "VALUE"},
+        auto_remove=True,
+    )
+
+    mocker_normalize.assert_called_once_with("test-image")
+    assert config["name"] == "test-image"
+    assert config["image"] == "test-image:latest"
+    assert config["command"] == ["echo", "hello"]
+    assert config["environment"] == {"KEY": "VALUE"}
+    assert config["mounts"] == [
+        {
+            "type": "bind",
+            "source": f"/source/{organization_id}/{workflow_id}/{execution_id}",
+            "target": "/data",
+        }
+    ]
+
+
+def test_get_container_run_config_without_mount(docker_client, mocker):
+    """Test the get_container_run_config method."""
+    os.environ[Env.WORKFLOW_DATA_DIR] = "/source"
+    command = ["echo", "hello"]
+
+    mocker.patch.object(docker_client, "_Client__image_exists", return_value=True)
+    mocker_normalize = mocker.patch(
+        f"{DOCKER_MODULE}.ContainerClientHelper.normalize_container_name",
+        return_value="test-image",
+    )
+    config = docker_client.get_container_run_config(
+        command,
+        None,
+        None,
+        None,
+        auto_remove=True,
+    )
+
+    mocker_normalize.assert_called_once_with("test-image")
+    assert config["name"] == "test-image"
+    assert config["image"] == "test-image:latest"
+    assert config["command"] == ["echo", "hello"]
+    assert config["environment"] == {}
+    assert config["mounts"] == []
+
+
+def test_run_container(docker_client, mocker):
+    """Test the run_container method."""
+    # Patch the client object to control its behavior
+    mock_client = mocker.patch.object(docker_client, "client")
+
+    config = {
+        "name": "test-image",
+        "image": "test-image:latest",
+        "command": ["echo", "hello"],
+        "detach": True,
+        "stream": True,
+        "auto_remove": True,
+        "environment": {"KEY": "VALUE"},
+        "stderr": True,
+        "stdout": True,
+        "network": "",
+        "mounts": [],
+    }
+
+    assert isinstance(docker_client.run_container(config), DockerContainer)
+    mock_client.containers.run.assert_called_once_with(**config)
+
+
+if __name__ == "__main__":
+    pytest.main()

worker/src/unstract/worker/worker.py

@@ -6,7 +6,7 @@
 
 from dotenv import load_dotenv
 from flask import Flask
-from unstract.worker.clients.helper import get_container_client
+from unstract.worker.clients.helper import ContainerClientHelper
 from unstract.worker.clients.interface import (
     ContainerClientInterface,
     ContainerInterface,
@@ -18,7 +18,7 @@
 
 load_dotenv()
 # Loads the container clinet class.
-client_class = get_container_client()
+client_class = ContainerClientHelper.get_container_client()
 
 
 class UnstractWorker: