From 22a906dbf73f78d8d9bee533e1dbd1bd0cf9dc9c Mon Sep 17 00:00:00 2001
From: Tomer Nosrati <tomer.nosrati@kcg.tech>
Date: Mon, 27 Nov 2023 18:16:12 +0200
Subject: [PATCH] Allow disabling clusters by returning None when overriding
 fixtures: celery_broker_cluster, celery_backend_cluster

---
 src/pytest_celery/api/setup.py                | 30 +++++++++++++------
 src/pytest_celery/fixtures/backend.py         |  6 ++--
 src/pytest_celery/fixtures/broker.py          |  6 ++--
 src/pytest_celery/vendors/worker/container.py |  4 +++
 4 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/src/pytest_celery/api/setup.py b/src/pytest_celery/api/setup.py
index 855bae05e..8de969b16 100644
--- a/src/pytest_celery/api/setup.py
+++ b/src/pytest_celery/api/setup.py
@@ -66,12 +66,14 @@ def config(cls, celery_worker_cluster_config: dict) -> dict:
         if not celery_worker_cluster_config:
             raise ValueError("celery_worker_cluster_config is empty")
 
-        celery_broker_cluster_config: dict = celery_worker_cluster_config["celery_broker_cluster_config"]
-        celery_backend_cluster_config: dict = celery_worker_cluster_config["celery_backend_cluster_config"]
-        return {
-            "broker_url": ";".join(celery_broker_cluster_config["local_urls"]),
-            "result_backend": ";".join(celery_backend_cluster_config["local_urls"]),
-        }
+        celery_broker_cluster_config: dict = celery_worker_cluster_config.get("celery_broker_cluster_config", {})
+        celery_backend_cluster_config: dict = celery_worker_cluster_config.get("celery_backend_cluster_config", {})
+        config = {}
+        if celery_broker_cluster_config:
+            config["broker_url"] = ";".join(celery_broker_cluster_config["local_urls"])
+        if celery_backend_cluster_config:
+            config["result_backend"] = ";".join(celery_backend_cluster_config["local_urls"])
+        return config
 
     @classmethod
     def update_app_config(cls, app: Celery) -> None:
@@ -110,8 +112,12 @@ def ready(self, ping: bool = False, control: bool = False, docker: bool = True)
         ready = True
 
         if docker and ready:
-            ready = all([self.broker_cluster.ready(), self.backend_cluster.ready()])
-            ready = ready and self.worker_cluster.ready()
+            if self.broker_cluster:
+                ready = ready and self.broker_cluster.ready()
+            if self.backend_cluster:
+                ready = ready and self.backend_cluster.ready()
+            if self.worker_cluster:
+                ready = ready and self.worker_cluster.ready()
 
         if control and ready:
             r = self.app.control.ping()
@@ -125,7 +131,13 @@ def ready(self, ping: bool = False, control: bool = False, docker: bool = True)
                 ready = ready and res.get(timeout=RESULT_TIMEOUT) == "pong"
 
         # Set app for all nodes
-        nodes = self.broker_cluster.nodes + self.backend_cluster.nodes
+        nodes: tuple = tuple()
+        if self.broker_cluster:
+            nodes += self.broker_cluster.nodes
+        if self.backend_cluster:
+            nodes += self.backend_cluster.nodes
+        if self.worker_cluster:
+            nodes += self.worker_cluster.nodes
         for node in nodes:
             node._app = self.app
 
diff --git a/src/pytest_celery/fixtures/backend.py b/src/pytest_celery/fixtures/backend.py
index d6984b916..9e57541a0 100644
--- a/src/pytest_celery/fixtures/backend.py
+++ b/src/pytest_celery/fixtures/backend.py
@@ -1,5 +1,7 @@
 # mypy: disable-error-code="misc"
 
+from typing import Any
+
 import pytest
 
 from pytest_celery.api.backend import CeleryBackendCluster
@@ -23,10 +25,10 @@ def celery_backend_cluster(celery_backend: CeleryTestBackend) -> CeleryBackendCl
 
 
 @pytest.fixture
-def celery_backend_cluster_config(request: pytest.FixtureRequest) -> dict:
+def celery_backend_cluster_config(request: pytest.FixtureRequest) -> dict[Any, Any] | None:
     try:
         use_default_config = pytest.fail.Exception
         cluster: CeleryBackendCluster = request.getfixturevalue(CELERY_BACKEND_CLUSTER)
-        return cluster.config()
+        return cluster.config() if cluster else None
     except use_default_config:
         return CeleryBackendCluster.default_config()
diff --git a/src/pytest_celery/fixtures/broker.py b/src/pytest_celery/fixtures/broker.py
index 86710abf0..01a414dbd 100644
--- a/src/pytest_celery/fixtures/broker.py
+++ b/src/pytest_celery/fixtures/broker.py
@@ -1,5 +1,7 @@
 # mypy: disable-error-code="misc"
 
+from typing import Any
+
 import pytest
 
 from pytest_celery.api.broker import CeleryBrokerCluster
@@ -23,10 +25,10 @@ def celery_broker_cluster(celery_broker: CeleryTestBroker) -> CeleryBrokerCluste
 
 
 @pytest.fixture
-def celery_broker_cluster_config(request: pytest.FixtureRequest) -> dict:
+def celery_broker_cluster_config(request: pytest.FixtureRequest) -> dict[Any, Any] | None:
     try:
         use_default_config = pytest.fail.Exception
         cluster: CeleryBrokerCluster = request.getfixturevalue(CELERY_BROKER_CLUSTER)
-        return cluster.config()
+        return cluster.config() if cluster else None
     except use_default_config:
         return CeleryBrokerCluster.default_config()
diff --git a/src/pytest_celery/vendors/worker/container.py b/src/pytest_celery/vendors/worker/container.py
index b15745ae6..fef4b8994 100644
--- a/src/pytest_celery/vendors/worker/container.py
+++ b/src/pytest_celery/vendors/worker/container.py
@@ -60,8 +60,12 @@ def env(cls, celery_worker_cluster_config: dict) -> dict:
         env = {}
         if celery_broker_cluster_config:
             env["CELERY_BROKER_URL"] = ";".join(celery_broker_cluster_config["urls"])
+        else:
+            del DEFAULT_WORKER_ENV["CELERY_BROKER_URL"]
         if celery_backend_cluster_config:
             env["CELERY_RESULT_BACKEND"] = ";".join(celery_backend_cluster_config["urls"])
+        else:
+            del DEFAULT_WORKER_ENV["CELERY_RESULT_BACKEND"]
         return {**DEFAULT_WORKER_ENV, **env}
 
     @classmethod