Move KindLoad to a condition instead of a wrapper (#21292) (#21307)

* Move KindLoad to a condition instead of a wrapper

* Add changelog

* Clean up typing

(cherry picked from commit 586c0a2)

Co-authored-by: Juanpe Araque <[email protected]>

Author: datadog-agent-integrations-bot[bot]

datadog_checks_dev/changelog.d/21292.fixed

@@ -0,0 +1 @@
+Move KindLoad to a built-in condition instead of a wrapper

datadog_checks_dev/datadog_checks/dev/kind.py

@@ -4,8 +4,9 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
+from dataclasses import dataclass
 from shutil import which
-from typing import TYPE_CHECKING
+from typing import Any, Callable, ContextManager, Protocol, runtime_checkable
 
 import pytest
 
@@ -15,59 +16,49 @@
 from .subprocess import run_command
 from .utils import get_active_env, get_current_check_name
 
-if TYPE_CHECKING:
-    from contextlib import AbstractContextManager
-    from typing import Self
 
+def _setup_conditions(conditions: list[Callable[[], Any]] | None, cluster_config: ClusterConfig):
+    if not conditions:
+        return
 
-def _setup_wrappers(wrappers: list[AbstractContextManager] | None, cluster_name: str):
-    """Set up wrappers with cluster-specific configuration.
+    for condition in conditions:
+        if isinstance(condition, ClusterCondition):
+            condition.add_cluster_info(cluster_config)
 
-    :param wrappers: List of wrapper instances to configure
-    :param cluster_name: The name of the Kind cluster
-    """
-    if not wrappers:
-        return
 
-    for wrapper in wrappers:
-        match wrapper:
-            case KindLoad():
-                wrapper.cluster_name = cluster_name
-            case _:
-                # No special setup needed for other wrapper types
-                pass
+@dataclass
+class ClusterConfig:
+    cluster_name: str
+
+
+@runtime_checkable
+class ClusterCondition(Protocol):
+    def add_cluster_info(self, cluster_config: ClusterConfig): ...
 
 
 @contextmanager
 def kind_run(
-    sleep=None,
-    endpoints=None,
-    conditions=None,
-    env_vars=None,
-    wrappers=None,
-    kind_config=None,
-    attempts=None,
-    attempts_wait=1,
+    sleep: float | None = None,
+    endpoints: str | list[str] | None = None,
+    conditions: list[Callable[[], Any]] | None = None,
+    env_vars: dict[str, str] | None = None,
+    wrappers: list[ContextManager] | None = None,
+    kind_config: str | None = None,
+    attempts: int | None = None,
+    attempts_wait: int = 1,
 ):
     """
     This utility provides a convenient way to safely set up and tear down Kind environments.
 
     :param sleep: Number of seconds to wait before yielding.
-    :type sleep: ``float``
     :param endpoints: Endpoints to verify access for before yielding. Shorthand for adding
                       ``conditions.CheckEndpoints(endpoints)`` to the ``conditions`` argument.
-    :type endpoints: ``list`` of ``str``, or a single ``str``
     :param conditions: A list of callable objects that will be executed before yielding to check for errors.
-    :type conditions: ``callable``
     :param env_vars: A dictionary to update ``os.environ`` with during execution.
-    :type env_vars: ``dict``
     :param wrappers: A list of context managers to use during execution.
     :param kind_config: A path to a yaml file that contains the configuration for creating the kind cluster.
-    :type kind_config: ``str``
     :param attempts: Number of attempts to run `up` and the `conditions` successfully. Defaults to 2 in CI.
-    :type attempts: ``int``
     :param attempts_wait: Time to wait between attempts.
-    :type attempts_wait: ``int``
     """
     if not which('kind'):
         pytest.skip('Kind not available')
@@ -88,8 +79,7 @@ def kind_run(
             set_up = KindUp(cluster_name, kind_config)
             tear_down = KindDown(cluster_name)
 
-            # Set up wrappers with cluster-specific configuration
-            _setup_wrappers(wrappers, cluster_name)
+            _setup_conditions(conditions, ClusterConfig(cluster_name))
 
             with environment_run(
                 up=set_up,
@@ -135,14 +125,14 @@ def __call__(self):
         run_command(['kind', 'delete', 'cluster', '--name', self.cluster_name], check=True)
 
 
-class KindLoad:
-    """Context manager for loading Docker images into a Kind cluster.
+class KindLoad(LazyFunction):
+    """Condition for loading Docker images into a Kind cluster.
 
-    This context manager should be passed to the wrappers argument in environment_run
+    This condition should be passed to the conditions argument in environment_run
     to load images into the Kind cluster after it's created.
 
     Example:
-        with kind_run(wrappers=[KindLoad("my-image:latest")]):
+        with kind_run(conditions=[KindLoad("my-image:latest")]):
             # The image is now loaded in the kind cluster
             pass
     """
@@ -151,14 +141,12 @@ def __init__(self, image: str):
         self.image = image
         self.cluster_name: str | None = None
 
-    def __enter__(self) -> Self:
+    def __call__(self):
         if self.cluster_name is None:
-            raise RuntimeError("cluster_name must be set before entering KindLoad context")
+            raise RuntimeError("cluster_name must be set before calling KindLoad")
 
         load_cmd = ['kind', 'load', 'docker-image', self.image, '--name', self.cluster_name]
         run_command(load_cmd, check=True)
-        return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Exit the context manager (no cleanup needed for image loading)."""
-        pass
+    def add_cluster_info(self, cluster_config: ClusterConfig):
+        self.cluster_name = cluster_config.cluster_name

datadog_checks_dev/tests/test_kind.py

@@ -61,23 +61,20 @@ def test_retry_condition_failed_only_on_first_run(self):
 
 
 class TestKindLoad:
-    def test_kind_load_context_manager_without_cluster_name(self):
+    def test_kind_load_without_cluster_name(self):
         kind_load = KindLoad("test-image:latest")
 
-        with pytest.raises(RuntimeError, match="cluster_name must be set before entering KindLoad context"):
-            with kind_load:
-                pass
+        with pytest.raises(RuntimeError, match="cluster_name must be set before calling KindLoad"):
+            kind_load()
 
     @patch('datadog_checks.dev.kind.run_command')
-    def test_kind_load_context_manager_with_cluster_name(self, mock_run_command):
-        """Test that KindLoad calls the correct kind load command when cluster_name is set."""
+    def test_kind_load_with_cluster_name(self, mock_run_command):
         image = "test-image:latest"
         cluster_name = "test-cluster"
         kind_load = KindLoad(image)
         kind_load.cluster_name = cluster_name
 
-        with kind_load as ctx:
-            assert ctx is kind_load
+        kind_load()
 
         mock_run_command.assert_called_once_with(
             ['kind', 'load', 'docker-image', image, '--name', cluster_name], check=True
@@ -86,7 +83,6 @@ def test_kind_load_context_manager_with_cluster_name(self, mock_run_command):
     @not_windows_ci
     @patch('datadog_checks.dev.kind.run_command')
     def test_kind_load_integration_with_kind_run(self, mock_run_command):
-        """Test that KindLoad integrates correctly with kind_run."""
         image = "test-image:latest"
         kind_load = KindLoad(image)
 
@@ -100,7 +96,7 @@ def test_kind_load_integration_with_kind_run(self, mock_run_command):
             mock_down_instance = MagicMock()
             mock_kind_down.return_value = mock_down_instance
 
-            with kind_run(wrappers=[kind_load]):
+            with kind_run(conditions=[kind_load]):
                 # Verify that cluster_name was set on the KindLoad instance
                 assert kind_load.cluster_name is not None
                 assert kind_load.cluster_name.startswith('cluster-')