Refactor to use ProxiedDockerContainerExtension with better abstractions

- Add env_vars and health_check_fn parameters to ProxiedDockerContainerExtension
- Refactor ParadeDbExtension to use shared base class
- Improve integration tests to use ProxiedDockerContainerExtension
- Remove weak test assertions and fix flaky tests

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: whummer

paradedb/localstack_paradedb/extension.py

@@ -4,9 +4,6 @@
 
 from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
 from localstack.extensions.api import http
-from localstack.utils.docker_utils import DOCKER_CLIENT
-from localstack.utils.container_utils.container_client import PortMappings
-from localstack.utils.sync import retry
 from werkzeug.datastructures import Headers
 
 LOG = logging.getLogger(__name__)
@@ -39,22 +36,28 @@ def __init__(self):
         postgres_db = os.environ.get(ENV_POSTGRES_DB, DEFAULT_POSTGRES_DB)
         postgres_port = int(os.environ.get(ENV_POSTGRES_PORT, DEFAULT_POSTGRES_PORT))
 
+        # Store configuration for connection info
+        self.postgres_user = postgres_user
+        self.postgres_password = postgres_password
+        self.postgres_db = postgres_db
+        self.postgres_port = postgres_port
+
         # Environment variables to pass to the container
-        self.env_vars = {
+        env_vars = {
             "POSTGRES_USER": postgres_user,
             "POSTGRES_PASSWORD": postgres_password,
             "POSTGRES_DB": postgres_db,
         }
 
-        # Store configuration for connection info
-        self.postgres_user = postgres_user
-        self.postgres_password = postgres_password
-        self.postgres_db = postgres_db
-        self.postgres_port = postgres_port
+        def _tcp_health_check():
+            """Check if PostgreSQL port is accepting connections."""
+            self._check_tcp_port(self.container_host, self.postgres_port)
 
         super().__init__(
             image_name=self.DOCKER_IMAGE,
             container_ports=[postgres_port],
+            env_vars=env_vars,
+            health_check_fn=_tcp_health_check,
         )
 
     def should_proxy_request(self, headers: Headers) -> bool:
@@ -71,45 +74,6 @@ def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
         """
         self.start_container()
 
-    def start_container(self) -> None:
-        """Override to add env_vars support and database-specific health checking."""
-        LOG.debug("Starting extension container %s", self.container_name)
-
-        port_mapping = PortMappings()
-        for port in self.container_ports:
-            port_mapping.add(port)
-
-        try:
-            DOCKER_CLIENT.run_container(
-                self.image_name,
-                detach=True,
-                remove=True,
-                name=self.container_name,
-                ports=port_mapping,
-                env_vars=self.env_vars,
-            )
-        except Exception as e:
-            LOG.debug("Failed to start container %s: %s", self.container_name, e)
-            raise
-
-        def _check_health():
-            """Check if PostgreSQL port is accepting connections."""
-            self._check_tcp_port(self.container_host, self.postgres_port)
-
-        try:
-            retry(_check_health, retries=60, sleep=1)
-        except Exception as e:
-            LOG.info("Failed to connect to container %s: %s", self.container_name, e)
-            self._remove_container()
-            raise
-
-        LOG.info(
-            "Successfully started extension container %s on %s:%s",
-            self.container_name,
-            self.container_host,
-            self.postgres_port,
-        )
-
     def _check_tcp_port(self, host: str, port: int, timeout: float = 2.0) -> None:
         """Check if a TCP port is accepting connections."""
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

utils/localstack_extensions/utils/docker.py

@@ -47,6 +47,13 @@ class ProxiedDockerContainerExtension(Extension):
     """Optional path on which to expose the container endpoints."""
     command: list[str] | None
     """Optional command (and flags) to execute in the container."""
+    env_vars: dict[str, str] | None
+    """Optional environment variables to pass to the container."""
+    health_check_fn: Callable[[], None] | None
+    """
+    Optional custom health check function. If not provided, defaults to HTTP GET on main_port.
+    The function should raise an exception if the health check fails.
+    """
 
     request_to_port_router: Callable[[Request], int] | None
     """Callable that returns the target port for a given request, for routing purposes"""
@@ -60,6 +67,8 @@ def __init__(
         host: str | None = None,
         path: str | None = None,
         command: list[str] | None = None,
+        env_vars: dict[str, str] | None = None,
+        health_check_fn: Callable[[], None] | None = None,
         request_to_port_router: Callable[[Request], int] | None = None,
         http2_ports: list[int] | None = None,
     ):
@@ -71,6 +80,8 @@ def __init__(
         self.path = path
         self.container_name = re.sub(r"\W", "-", f"ls-ext-{self.name}")
         self.command = command
+        self.env_vars = env_vars
+        self.health_check_fn = health_check_fn
         self.request_to_port_router = request_to_port_router
         self.http2_ports = http2_ports
         self.main_port = self.container_ports[0]
@@ -113,6 +124,8 @@ def start_container(self) -> None:
         kwargs = {}
         if self.command:
             kwargs["command"] = self.command
+        if self.env_vars:
+            kwargs["env_vars"] = self.env_vars
 
         try:
             DOCKER_CLIENT.run_container(
@@ -129,20 +142,23 @@ def start_container(self) -> None:
             if not is_env_true(f"{self.name.upper().replace('-', '_')}_DEV_MODE"):
                 raise
 
-        def _ping_endpoint():
-            # TODO: allow defining a custom healthcheck endpoint ...
-            response = requests.get(f"http://{self.container_host}:{self.main_port}/")
-            assert response.ok
+        # Use custom health check if provided, otherwise default to HTTP GET
+        health_check = self.health_check_fn or self._default_health_check
 
         try:
-            retry(_ping_endpoint, retries=40, sleep=1)
+            retry(health_check, retries=60, sleep=1)
         except Exception as e:
             LOG.info("Failed to connect to container %s: %s", self.container_name, e)
             self._remove_container()
             raise
 
         LOG.debug("Successfully started extension container %s", self.container_name)
 
+    def _default_health_check(self) -> None:
+        """Default health check: HTTP GET request to the main port."""
+        response = requests.get(f"http://{self.container_host}:{self.main_port}/")
+        assert response.ok
+
     def _remove_container(self):
         LOG.debug("Stopping extension container %s", self.container_name)
         DOCKER_CLIENT.remove_container(

utils/tests/integration/conftest.py

@@ -3,99 +3,93 @@
 
 Provides fixtures for running tests against the grpcbin Docker container.
 grpcbin is a neutral gRPC test service that supports various RPC types.
+
+Uses ProxiedDockerContainerExtension to manage the grpcbin container,
+providing realistic test coverage of the Docker container management infrastructure.
 """
 
-import time
+import socket
 import pytest
 
-from localstack.utils.container_utils.container_client import PortMappings
-from localstack.utils.docker_utils import DOCKER_CLIENT
-from localstack.utils.net import wait_for_port_open
+from werkzeug.datastructures import Headers
+from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
 
 
 GRPCBIN_IMAGE = "moul/grpcbin"
 GRPCBIN_INSECURE_PORT = 9000  # HTTP/2 without TLS
 GRPCBIN_SECURE_PORT = 9001  # HTTP/2 with TLS
 
 
-@pytest.fixture(scope="session")
-def grpcbin_container():
+class GrpcbinExtension(ProxiedDockerContainerExtension):
+    """
+    Test extension for grpcbin that uses ProxiedDockerContainerExtension.
+
+    This extension demonstrates using ProxiedDockerContainerExtension for
+    a gRPC/HTTP2 service. While grpcbin doesn't use the HTTP gateway routing
+    (it's accessed via direct TCP), this tests the Docker container management
+    capabilities of ProxiedDockerContainerExtension.
     """
-    Start a grpcbin Docker container for testing.
 
-    The container exposes:
-    - Port 9000: Insecure gRPC (HTTP/2 without TLS)
-    - Port 9001: Secure gRPC (HTTP/2 with TLS)
+    name = "grpcbin-test"
+
+    def __init__(self):
+        def _tcp_health_check():
+            """Check if grpcbin insecure port is accepting TCP connections."""
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.settimeout(2.0)
+            try:
+                # Use container_host from the parent class
+                sock.connect((self.container_host, GRPCBIN_INSECURE_PORT))
+                sock.close()
+            except (socket.timeout, socket.error) as e:
+                raise AssertionError(f"Port {GRPCBIN_INSECURE_PORT} not ready: {e}")
+
+        super().__init__(
+            image_name=GRPCBIN_IMAGE,
+            container_ports=[GRPCBIN_INSECURE_PORT, GRPCBIN_SECURE_PORT],
+            health_check_fn=_tcp_health_check,
+        )
+
+    def should_proxy_request(self, headers: Headers) -> bool:
+        """
+        gRPC services use direct TCP connections, not HTTP gateway routing.
+        This method is not used in these tests but is required by the base class.
+        """
+        return False
+
 
-    The container is automatically removed after tests complete.
+@pytest.fixture(scope="session")
+def grpcbin_extension():
     """
-    container_name = "pytest-grpcbin"
-
-    # Remove any existing container with the same name
-    try:
-        DOCKER_CLIENT.remove_container(container_name)
-    except Exception:
-        pass  # Container may not exist
-
-    # Configure port mappings
-    ports = PortMappings()
-    ports.add(GRPCBIN_INSECURE_PORT)
-    ports.add(GRPCBIN_SECURE_PORT)
-
-    # Start the container
-    stdout, _ = DOCKER_CLIENT.run_container(
-        image_name=GRPCBIN_IMAGE,
-        name=container_name,
-        detach=True,
-        remove=True,
-        ports=ports,
-    )
-    container_id = stdout.decode().strip()
-
-    # Wait for the insecure port to be ready with enough retries
-    try:
-        wait_for_port_open(GRPCBIN_INSECURE_PORT, retries=60, sleep_time=0.5)
-    except Exception:
-        # Clean up and fail
-        try:
-            DOCKER_CLIENT.remove_container(container_name)
-        except Exception:
-            pass
-        pytest.fail(f"grpcbin port {GRPCBIN_INSECURE_PORT} did not become available")
-
-    # Brief delay to allow grpcbin HTTP/2 server to fully initialize
-    # (port open doesn't guarantee HTTP/2 readiness)
-    time.sleep(0.5)
-
-    # Provide connection info to tests
-    yield {
-        "container_id": container_id,
-        "container_name": container_name,
-        "host": "localhost",
-        "insecure_port": GRPCBIN_INSECURE_PORT,
-        "secure_port": GRPCBIN_SECURE_PORT,
-    }
-
-    # Cleanup: stop and remove the container
-    try:
-        DOCKER_CLIENT.remove_container(container_name)
-    except Exception:
-        pass
+    Start grpcbin using ProxiedDockerContainerExtension.
+
+    This tests the Docker container management capabilities while providing
+    a realistic gRPC/HTTP2 test service for integration tests.
+    """
+    extension = GrpcbinExtension()
+
+    # Start the container using the extension infrastructure
+    extension.start_container()
+
+    yield extension
+
+    # Cleanup
+    extension.on_platform_shutdown()
 
 
 @pytest.fixture
-def grpcbin_host(grpcbin_container):
+def grpcbin_host(grpcbin_extension):
     """Return the host address for the grpcbin container."""
-    return grpcbin_container["host"]
+    return grpcbin_extension.container_host
 
 
 @pytest.fixture
-def grpcbin_insecure_port(grpcbin_container):
+def grpcbin_insecure_port(grpcbin_extension):
     """Return the insecure (HTTP/2 without TLS) port for grpcbin."""
-    return grpcbin_container["insecure_port"]
+    return GRPCBIN_INSECURE_PORT
 
 
 @pytest.fixture
-def grpcbin_secure_port(grpcbin_container):
+def grpcbin_secure_port(grpcbin_extension):
     """Return the secure (HTTP/2 with TLS) port for grpcbin."""
-    return grpcbin_container["secure_port"]
+    return GRPCBIN_SECURE_PORT

utils/tests/integration/test_extension_integration.py

@@ -0,0 +1,49 @@
+"""
+Integration tests for ProxiedDockerContainerExtension with grpcbin.
+
+These tests verify that ProxiedDockerContainerExtension properly manages
+Docker containers in a realistic scenario, using grpcbin as a test service.
+"""
+
+import socket
+
+
+class TestProxiedDockerContainerExtension:
+    """Tests for ProxiedDockerContainerExtension using the GrpcbinExtension."""
+
+    def test_extension_starts_container(self, grpcbin_extension):
+        """Test that the extension successfully starts the Docker container."""
+        assert grpcbin_extension.container_name == "ls-ext-grpcbin-test"
+        assert grpcbin_extension.image_name == "moul/grpcbin"
+        assert len(grpcbin_extension.container_ports) == 2
+
+    def test_extension_container_host_is_accessible(self, grpcbin_extension):
+        """Test that the container_host is set and accessible."""
+        assert grpcbin_extension.container_host is not None
+        # container_host should be localhost, localhost.localstack.cloud, or a docker bridge IP
+        assert grpcbin_extension.container_host in ("localhost", "127.0.0.1", "localhost.localstack.cloud") or \
+               grpcbin_extension.container_host.startswith("172.")
+
+    def test_extension_ports_are_reachable(self, grpcbin_host, grpcbin_insecure_port):
+        """Test that the extension's ports are reachable via TCP."""
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(2.0)
+        try:
+            sock.connect((grpcbin_host, grpcbin_insecure_port))
+            sock.close()
+            # Connection successful
+        except (socket.timeout, socket.error) as e:
+            raise AssertionError(f"Could not connect to grpcbin port: {e}")
+
+    def test_extension_implements_required_methods(self, grpcbin_extension):
+        """Test that the extension properly implements the required abstract methods."""
+        from werkzeug.datastructures import Headers
+
+        # should_proxy_request should be callable
+        result = grpcbin_extension.should_proxy_request(Headers())
+        assert result is False, "gRPC services should not proxy through HTTP gateway"
+
+    def test_multiple_ports_configured(self, grpcbin_extension):
+        """Test that the extension properly handles multiple ports."""
+        assert 9000 in grpcbin_extension.container_ports  # insecure port
+        assert 9001 in grpcbin_extension.container_ports  # secure port