SNOW-2912540: add mitmproxy capture fixture base

setup.py

@@ -73,6 +73,7 @@
     "psutil",  # testing for telemetry
     "lxml",  # used in XML reader unit tests
     "pyarrow",  # used in dataframe reader tests
+    "mitmproxy>=11,<13",  # wire-level telemetry capture in tests/mitmproxy_client.py
 ]
 MODIN_DEVELOPMENT_REQUIREMENTS = [
     # Snowpark pandas 3rd party library testing. Cap the scipy version because

tests/integ/conftest.py

@@ -19,6 +19,7 @@
     setup_full_ast_validation_mode,
 )
 from tests.integ.session_parameters import set_up_test_session_parameters
+from tests.mitmproxy_client import MitmproxyClient
 from tests.parameters import CONNECTION_PARAMETERS
 from tests.utils import (
     TEST_SCHEMA,
@@ -407,6 +408,52 @@ def session(
         session.close()
 
 
+@pytest.fixture(scope="session")
+def _mitmproxy_session():
+    """Start one mitmdump process per xdist worker and reuse it across tests."""
+    client = MitmproxyClient().start()
+    try:
+        yield client
+    finally:
+        client.stop()
+
+
+@pytest.fixture
+def mitmproxy(_mitmproxy_session):
+    """Per-test mitmproxy handle backed by a session-scoped mitmdump process.
+
+    Captured requests are cleared before each test; the process itself stays
+    up, saving the mitmdump startup cost per test.
+    """
+    _mitmproxy_session.reset()
+    return _mitmproxy_session
+
+
+@pytest.fixture
+def mitmproxy_session(db_parameters, local_testing_mode, mitmproxy, monkeypatch):
+    """A Session whose traffic is routed through the mitmproxy fixture.
+
+    Login/query/telemetry all reach the real account untouched -- the proxy
+    only observes traffic in transit, so no backend response is faked here.
+    """
+    if local_testing_mode:
+        pytest.skip("mitmproxy capture requires a real network connection")
+    monkeypatch.setenv("REQUESTS_CA_BUNDLE", str(mitmproxy.ca_cert_path))
+    params = dict(db_parameters)
+    params["proxy_host"] = mitmproxy.proxy_host
+    params["proxy_port"] = mitmproxy.proxy_port
+    # mitmproxy re-signs TLS traffic with its own CA-issued leaf cert, which
+    # has no real OCSP/CRL revocation record -- disable revocation checking
+    # rather than relying on ocsp_fail_open, since mitmproxy's cert triggers
+    # a hard "revoked" failure, not a soft "responder unreachable" one.
+    params["disable_ocsp_checks"] = True
+    proxied_session = Session.builder.configs(params).create()
+    try:
+        yield proxied_session
+    finally:
+        proxied_session.close()
+
+
 @pytest.fixture(scope="function")
 def profiler_session(
     db_parameters,

tests/integ/test_mitmproxy_fixture.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+"""Self-test proving the mitmproxy capture fixture works end-to-end.
+
+Exists so a break in the fixture itself (proxy routing, CA trust, capture)
+fails loudly here instead of silently as an empty-results false pass in the
+telemetry tests that depend on it.
+"""
+
+
+def test_mitmproxy_captures_login_request(mitmproxy_session, mitmproxy):
+    # mitmproxy_session already completed login by the time this test body
+    # runs, so the request should already be captured.
+    requests = mitmproxy.wait_for_requests(
+        r"/session/v1/login-request", min_count=1, timeout=10.0
+    )
+    assert (
+        len(requests) >= 1
+    ), "Expected the session's login request to be captured by the proxy"

tests/mitmproxy_addon.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+"""mitmdump addon that records every request mitmproxy sees.
+
+Loaded via `mitmdump -s tests/mitmproxy_addon.py`. Appends one JSON line per
+request to the file at MITMPROXY_CAPTURE_OUTPUT_PATH; MitmproxyClient reads
+that file to answer get_requests()/wait_for_requests(). Traffic is otherwise
+untouched -- this addon never returns a response, so requests continue to
+their real destination.
+"""
+import json
+import os
+
+
+def request(flow) -> None:
+    output_path = os.environ["MITMPROXY_CAPTURE_OUTPUT_PATH"]
+    record = {
+        "method": flow.request.method,
+        "url": flow.request.pretty_url,
+        "headers": dict(flow.request.headers),
+        "body": flow.request.text,
+    }
+    with open(output_path, "a") as f:
+        f.write(json.dumps(record) + "\n")

tests/mitmproxy_client.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+"""Passive HTTPS capture for integration tests, via mitmproxy.
+
+Unlike a stub server, this never fabricates a response: it runs `mitmdump` as
+a real forward proxy so login/query/telemetry traffic reaches its actual
+destination untouched, while tests/mitmproxy_addon.py records every request
+mitmproxy decrypts to a JSON-lines file this client reads.
+"""
+from __future__ import annotations
+
+import json
+import os
+import re
+import shutil
+import socket
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Any
+
+
+class MitmproxyClient:
+    _ADDON_PATH = Path(__file__).parent / "mitmproxy_addon.py"
+
+    def __init__(self) -> None:
+        self.host = "127.0.0.1"
+        self.port: int | None = None
+        self._process: subprocess.Popen | None = None
+        fd, path = tempfile.mkstemp(prefix="mitmproxy_capture_", suffix=".jsonl")
+        os.close(fd)
+        self._output_path = Path(path)
+
+    def start(self) -> MitmproxyClient:
+        if shutil.which("mitmdump") is None:
+            raise RuntimeError(
+                "mitmdump not found on PATH. Install with: pip install mitmproxy"
+            )
+        self.port = self._find_free_port()
+        env = {**os.environ, "MITMPROXY_CAPTURE_OUTPUT_PATH": str(self._output_path)}
+        self._process = subprocess.Popen(
+            [
+                "mitmdump",
+                "--listen-host",
+                self.host,
+                "--listen-port",
+                str(self.port),
+                "-s",
+                str(self._ADDON_PATH),
+            ],
+            env=env,
+            # Redirect to DEVNULL: an unread stdout/stderr pipe can deadlock
+            # mitmdump's logging thread once its buffer fills.
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        self._wait_for_port()
+        self._wait_for_ca_cert()
+        return self
+
+    @staticmethod
+    def _find_free_port() -> int:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("127.0.0.1", 0))
+            return s.getsockname()[1]
+
+    def _wait_for_port(self, timeout: float = 30.0) -> None:
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            if self._process.poll() is not None:
+                raise RuntimeError(
+                    f"mitmdump exited early with code {self._process.returncode}"
+                )
+            try:
+                with socket.create_connection((self.host, self.port), timeout=0.5):
+                    return
+            except OSError:
+                time.sleep(0.2)
+        raise RuntimeError("mitmdump did not become ready in time")
+
+    def _wait_for_ca_cert(self, timeout: float = 10.0) -> None:
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            if self.ca_cert_path.exists():
+                return
+            time.sleep(0.2)
+        raise RuntimeError("mitmproxy CA cert was not generated in time")
+
+    @property
+    def ca_cert_path(self) -> Path:
+        return Path.home() / ".mitmproxy" / "mitmproxy-ca-cert.pem"
+
+    @property
+    def proxy_host(self) -> str:
+        return self.host
+
+    @property
+    def proxy_port(self) -> int:
+        return self.port
+
+    def reset(self) -> None:
+        self._output_path.write_text("")
+
+    def get_requests(self, url_path_pattern: str) -> list[dict[str, Any]]:
+        if not self._output_path.exists():
+            return []
+        matches = []
+        for line in self._output_path.read_text().splitlines():
+            if not line:
+                continue
+            entry = json.loads(line)
+            if re.search(url_path_pattern, entry["url"]):
+                matches.append(entry)
+        return matches
+
+    def wait_for_requests(
+        self,
+        url_path_pattern: str,
+        min_count: int = 1,
+        timeout: float = 2.0,
+        poll_interval: float = 0.1,
+    ) -> list[dict[str, Any]]:
+        """Poll until at least `min_count` requests matching the pattern arrive.
+
+        Useful for asserting on requests sent asynchronously (e.g. telemetry).
+        """
+        deadline = time.time() + timeout
+        result: list[dict[str, Any]] = []
+        while time.time() < deadline:
+            result = self.get_requests(url_path_pattern)
+            if len(result) >= min_count:
+                return result
+            time.sleep(poll_interval)
+        return result
+
+    def stop(self) -> None:
+        if self._process is not None:
+            self._process.terminate()
+            try:
+                self._process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                self._process.kill()
+                self._process.wait()
+            self._process = None
+        if self._output_path.exists():
+            self._output_path.unlink()
+
+    def __enter__(self) -> MitmproxyClient:
+        return self.start()
+
+    def __exit__(self, *exc_info: Any) -> None:
+        self.stop()