refactor utils tests to use embedded twisted server; minor polishing

Author: whummer

paradedb/localstack_paradedb/extension.py

@@ -2,7 +2,6 @@
 import socket
 
 from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
-from werkzeug.datastructures import Headers
 from localstack import config
 
 # Environment variables for configuration
@@ -83,13 +82,6 @@ def tcp_connection_matcher(self, data: bytes) -> bool:
 
         return False
 
-    def http2_request_matcher(self, headers: Headers) -> bool:
-        """
-        Define whether an HTTP2 request should be proxied based on request headers.
-        Not used for TCP connections - see tcp_connection_matcher instead.
-        """
-        return False
-
     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

@@ -1,6 +1,5 @@
 import re
 import logging
-from abc import abstractmethod
 from functools import cache
 from typing import Callable
 import requests
@@ -188,9 +187,14 @@ def _setup_tcp_protocol_routing(self):
             f"Registered TCP extension {self.name} -> {self.container_host}:{target_port} on gateway"
         )
 
-    @abstractmethod
     def http2_request_matcher(self, headers: Headers) -> bool:
-        """Define whether an HTTP2 request should be proxied, based on request headers."""
+        """
+        Define whether an HTTP2 request should be proxied, based on request headers.
+
+        Default implementation returns False (no HTTP2 proxying).
+        Override this method in subclasses that need HTTP2 proxying.
+        """
+        return False
 
     def on_platform_shutdown(self):
         self._remove_container()

utils/tests/integration/conftest.py

@@ -9,12 +9,18 @@
 """
 
 import socket
-import pytest
+import threading
+import time
 
+import pytest
 from hyperframe.frame import Frame
-from werkzeug.datastructures import Headers
-from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
+from localstack.utils.net import get_free_tcp_port
+from rolo import Router
+from rolo.gateway import Gateway
+from twisted.internet import reactor
+from twisted.web import server as twisted_server
 
+from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
 
 GRPCBIN_IMAGE = "moul/grpcbin"
 GRPCBIN_INSECURE_PORT = 9000  # HTTP/2 without TLS
@@ -53,51 +59,75 @@ def _tcp_health_check():
             image_name=GRPCBIN_IMAGE,
             container_ports=[GRPCBIN_INSECURE_PORT, GRPCBIN_SECURE_PORT],
             health_check_fn=_tcp_health_check,
+            tcp_ports=[GRPCBIN_INSECURE_PORT],  # Enable raw TCP proxying for gRPC/HTTP2
         )
 
-    def http2_request_matcher(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
+    def tcp_connection_matcher(self, data: bytes) -> bool:
+        """Detect HTTP/2 connection preface to route gRPC/HTTP2 traffic."""
+        # HTTP/2 connections start with the connection preface
+        if len(data) >= len(HTTP2_PREFACE):
+            return data.startswith(HTTP2_PREFACE)
+        # Also match if we have partial preface data (for early detection)
+        return len(data) > 0 and HTTP2_PREFACE.startswith(data)
 
 
 @pytest.fixture(scope="session")
-def grpcbin_extension():
+def grpcbin_extension_server():
     """
-    Start grpcbin using ProxiedDockerContainerExtension.
+    Start grpcbin using ProxiedDockerContainerExtension with a test gateway server.
 
-    This tests the Docker container management capabilities while providing
-    a realistic gRPC/HTTP2 test service for integration tests.
+    This tests the Docker container management and proxy capabilities by:
+    1. Starting the grpcbin container via the extension
+    2. Setting up a Gateway with the extension's routes and TCP patches
+    3. Serving the Gateway on a test port via Twisted
+    4. Returning server info for end-to-end testing
     """
     extension = GrpcbinExtension()
 
-    # Start the container using the extension infrastructure
-    extension.start_container()
+    # Create router and update with extension routes
+    # This will start the grpcbin container and apply TCP protocol patches
+    router = Router()
+    extension.update_gateway_routes(router)
 
-    yield extension
+    # Create a Gateway with proper TCP support
+    # The TCP patches are applied by update_gateway_routes above
+    gateway = Gateway(router)
 
-    # Cleanup
-    extension.on_platform_shutdown()
+    # Start gateway on a test port using Twisted
+    test_port = get_free_tcp_port()
+    site = twisted_server.Site(gateway)
+    listener = reactor.listenTCP(test_port, site)
 
+    # Run reactor in background thread
+    def run_reactor():
+        reactor.run(installSignalHandlers=False)
 
-@pytest.fixture
-def grpcbin_host(grpcbin_extension):
-    """Return the host address for the grpcbin container."""
-    return grpcbin_extension.container_host
+    reactor_thread = threading.Thread(target=run_reactor, daemon=True)
+    reactor_thread.start()
 
+    # Wait for reactor to start - not ideal, but should work as a simple solution
+    time.sleep(0.5)
 
-@pytest.fixture
-def grpcbin_insecure_port(grpcbin_extension):
-    """Return the insecure (HTTP/2 without TLS) port for grpcbin."""
-    return GRPCBIN_INSECURE_PORT
+    # Return server information for tests
+    server_info = {
+        "port": test_port,
+        "url": f"http://localhost:{test_port}",
+        "extension": extension,
+        "listener": listener,
+    }
 
+    yield server_info
 
-@pytest.fixture
-def grpcbin_secure_port(grpcbin_extension):
-    """Return the secure (HTTP/2 with TLS) port for grpcbin."""
-    return GRPCBIN_SECURE_PORT
+    # Cleanup
+    reactor.callFromThread(reactor.stop)
+    time.sleep(0.5)
+    extension.on_platform_shutdown()
+
+
+@pytest.fixture(scope="session")
+def grpcbin_extension(grpcbin_extension_server):
+    """Return the extension instance from the server fixture."""
+    return grpcbin_extension_server["extension"]
 
 
 def parse_server_frames(data: bytes) -> list:

utils/tests/integration/test_extension_integration.py

@@ -7,6 +7,8 @@
 
 import socket
 
+from werkzeug.datastructures import Headers
+
 
 class TestProxiedDockerContainerExtension:
     """Tests for ProxiedDockerContainerExtension using the GrpcbinExtension."""
@@ -27,21 +29,20 @@ def test_extension_container_host_is_accessible(self, grpcbin_extension):
             "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."""
+    def test_extension_ports_are_reachable(self, grpcbin_extension_server):
+        """Test that the gateway port is reachable via TCP."""
+        gateway_port = grpcbin_extension_server["port"]
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.settimeout(2.0)
         try:
-            sock.connect((grpcbin_host, grpcbin_insecure_port))
+            sock.connect(("localhost", gateway_port))
             sock.close()
             # Connection successful
         except (socket.timeout, socket.error) as e:
-            raise AssertionError(f"Could not connect to grpcbin port: {e}")
+            raise AssertionError(f"Could not connect to gateway 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
-
         # http2_request_matcher should be callable
         result = grpcbin_extension.http2_request_matcher(Headers())
         assert result is False, "gRPC services should not proxy through HTTP gateway"

utils/tests/integration/test_grpc_e2e.py

@@ -14,10 +14,11 @@
 class TestGrpcEndToEnd:
     """End-to-end tests making actual gRPC calls to grpcbin."""
 
-    def test_grpc_empty_call(self, grpcbin_host, grpcbin_insecure_port):
-        """Test making a gRPC call to grpcbin's Empty service."""
-        # Create a channel to grpcbin
-        channel = grpc.insecure_channel(f"{grpcbin_host}:{grpcbin_insecure_port}")
+    def test_grpc_empty_call(self, grpcbin_extension_server):
+        """Test making a gRPC call to grpcbin's Empty service via the gateway."""
+        # Create a channel to grpcbin through the gateway
+        gateway_port = grpcbin_extension_server["port"]
+        channel = grpc.insecure_channel(f"localhost:{gateway_port}")
 
         try:
             # Use grpc.channel_ready_future to verify connection
@@ -43,9 +44,10 @@ def test_grpc_empty_call(self, grpcbin_host, grpcbin_insecure_port):
         finally:
             channel.close()
 
-    def test_grpc_index_call(self, grpcbin_host, grpcbin_insecure_port):
+    def test_grpc_index_call(self, grpcbin_extension_server):
         """Test calling grpcbin's Index service which returns server info."""
-        channel = grpc.insecure_channel(f"{grpcbin_host}:{grpcbin_insecure_port}")
+        gateway_port = grpcbin_extension_server["port"]
+        channel = grpc.insecure_channel(f"localhost:{gateway_port}")
 
         try:
             # Verify channel is ready
@@ -68,9 +70,10 @@ def test_grpc_index_call(self, grpcbin_host, grpcbin_insecure_port):
         finally:
             channel.close()
 
-    def test_grpc_concurrent_calls(self, grpcbin_host, grpcbin_insecure_port):
+    def test_grpc_concurrent_calls(self, grpcbin_extension_server):
         """Test making multiple concurrent gRPC calls."""
-        channel = grpc.insecure_channel(f"{grpcbin_host}:{grpcbin_insecure_port}")
+        gateway_port = grpcbin_extension_server["port"]
+        channel = grpc.insecure_channel(f"localhost:{gateway_port}")
 
         try:
             # Verify channel is ready
@@ -97,9 +100,10 @@ def test_grpc_concurrent_calls(self, grpcbin_host, grpcbin_insecure_port):
         finally:
             channel.close()
 
-    def test_grpc_connection_reuse(self, grpcbin_host, grpcbin_insecure_port):
+    def test_grpc_connection_reuse(self, grpcbin_extension_server):
         """Test that a single gRPC channel can handle multiple sequential calls."""
-        channel = grpc.insecure_channel(f"{grpcbin_host}:{grpcbin_insecure_port}")
+        gateway_port = grpcbin_extension_server["port"]
+        channel = grpc.insecure_channel(f"localhost:{gateway_port}")
 
         try:
             # Verify channel is ready