Added psutil support to worker container (#114)

* Refactored src/pytest_celery/vendors/worker/Dockerfile: Split pip packages

* Refactored WorkerInitialContent.generate(): Split logic into private methods

* src/pytest_celery/vendors/worker/app.py -> src/pytest_celery/vendors/worker/content/app.py

* Added 'utils.py' to /app in the worker's container

* Added psutil package and CeleryTestWorker.get_running_processes_info() API

* Fixed unit tests

* Fixed error with type annotations in Python < 3.10

* Fixed mypy errors

* Reduced reruns for parallel CI runs to default (a.k.a tox.ini)

* Reduced reruns for integration & smoke CI runs from 10 -> 2 reruns

Author: Nusnus

.github/workflows/parallel-support.yml

@@ -66,9 +66,9 @@ jobs:
           poetry install --with dev,test,ci
 
       - name: Run tox for all environments in parallel
-        timeout-minutes: 20
+        timeout-minutes: 15
         run: |
-          tox -e xdist -- --reruns 10
+          tox -e xdist
 
   parallel:
     runs-on: ${{ matrix.os }}
@@ -112,6 +112,6 @@ jobs:
           poetry install --with dev,test,ci
 
       - name: Run tox for all environments in parallel
-        timeout-minutes: 20
+        timeout-minutes: 15
         run: |
-          tox -e parallel -- --reruns 10
+          tox -e parallel

.github/workflows/python-package.yml

@@ -123,7 +123,7 @@ jobs:
       - name: Run tox for "${{ matrix.python-version }}-integration"
         timeout-minutes: 15
         run: |
-          tox --verbose --verbose -e "${{ matrix.python-version }}-integration" -- -n auto --reruns 10 --rerun-except AssertionError
+          tox --verbose --verbose -e "${{ matrix.python-version }}-integration" -- -n auto --reruns 2 --rerun-except AssertionError
 
   Smoke:
     needs:
@@ -172,4 +172,4 @@ jobs:
       - name: Run tox for "${{ matrix.python-version }}-smoke"
         timeout-minutes: 15
         run: |
-          tox --verbose --verbose -e "${{ matrix.python-version }}-smoke" -- -n auto --reruns 10 --rerun-except AssertionError
+          tox --verbose --verbose -e "${{ matrix.python-version }}-smoke" -- -n auto --reruns 2 --rerun-except AssertionError

poetry.lock

@@ -71,6 +71,7 @@ celery = { version = "^5", extras = ["redis", "memcache", "pymemcache"] }
 retry = "^0.9.2"
 pytest-docker-tools = "^3.1.3"
 docker = "^7.0.0"
+psutil = "^5.9.7"
 
 [tool.poetry.group.dev]
 

pyproject.toml

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import json
+
 from celery import Celery
 
 from pytest_celery.api.base import CeleryTestCluster
@@ -34,6 +36,26 @@ def worker_queue(self) -> str:
     def hostname(self) -> str:
         return f"{self.worker_name}@{super().hostname()}"
 
+    def get_running_processes_info(
+        self,
+        columns: list[str] | None = None,
+        filters: dict[str, str] | None = None,
+    ) -> list[dict]:
+        exit_code, output = self.container.exec_run(
+            f'python -c "from utils import get_running_processes_info; print(get_running_processes_info({columns!r}))"'
+        )
+
+        if exit_code != 0:
+            raise RuntimeError(f"Failed to get processes info: {output}")
+
+        decoded_str = output.decode("utf-8")
+        output = json.loads(decoded_str)
+
+        if filters:
+            output = [item for item in output if all(item.get(key) == value for key, value in filters.items())]
+
+        return output
+
 
 class CeleryWorkerCluster(CeleryTestCluster):
     def __init__(self, *workers: tuple[CeleryTestWorker | CeleryTestContainer]) -> None:

src/pytest_celery/api/worker.py

@@ -19,13 +19,15 @@ ENV PYTHONUNBUFFERED=1
 ENV PYTHONDONTWRITEBYTECODE=1
 
 # Install Python dependencies
-RUN pip install --no-cache-dir --upgrade pip \
-    && pip install --no-cache-dir celery[redis,memcache,pymemcache]${CELERY_VERSION:+==$CELERY_VERSION}
+RUN pip install --no-cache-dir --upgrade \
+    pip \
+    celery[redis,memcache,pymemcache]${CELERY_VERSION:+==$CELERY_VERSION} \
+    psutil
 
 # The workdir must be /app
 WORKDIR /app
 
-COPY app.py app.py
+COPY content/ .
 
 # Switch to the test_user
 USER test_user

src/pytest_celery/vendors/worker/Dockerfile

@@ -39,10 +39,16 @@ def worker_queue(cls) -> str:
 
     @classmethod
     def app_module(cls) -> ModuleType:
-        from pytest_celery.vendors.worker import app
+        from pytest_celery.vendors.worker.content import app
 
         return app
 
+    @classmethod
+    def utils_module(cls) -> ModuleType:
+        from pytest_celery.vendors.worker.content import utils
+
+        return utils
+
     @classmethod
     def tasks_modules(cls) -> set:
         from pytest_celery.vendors.worker import tasks
@@ -88,15 +94,20 @@ def initial_content(
         worker_signals: set | None = None,
         worker_app: Celery | None = None,
         app_module: ModuleType | None = None,
+        utils_module: ModuleType | None = None,
     ) -> dict:
         if app_module is None:
             app_module = cls.app_module()
 
+        if utils_module is None:
+            utils_module = cls.utils_module()
+
         if worker_tasks is None:
             worker_tasks = cls.tasks_modules()
 
         content = WorkerInitialContent()
         content.set_app_module(app_module)
+        content.set_utils_module(utils_module)
         content.add_modules("tasks", worker_tasks)
         if worker_signals:
             content.add_modules("signals", worker_signals)

src/pytest_celery/vendors/worker/container.py

@@ -1,4 +1,5 @@
 """ Template for Celery worker application. """
+from __future__ import annotations
 
 import json
 import logging

src/pytest_celery/vendors/worker/content/__init__.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import json
+
+import psutil
+
+
+def get_running_processes_info(columns: list[str] | None = None) -> str:
+    if not columns:
+        columns = [
+            "pid",
+            "name",
+            "username",
+            "cmdline",
+            "cpu_percent",
+            "memory_percent",
+            "create_time",
+        ]
+    processes = [proc.info for proc in psutil.process_iter(columns)]
+    return json.dumps(processes)