SNOW-921045 Add wiremock tests support

.github/workflows/build_test.yml

@@ -138,6 +138,15 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Display Python version
         run: python -c "import sys; print(sys.version)"
+      - name: Set up Java
+        uses: actions/setup-java@v4 # for wiremock
+        with:
+          java-version: 11
+          distribution: 'temurin'
+          java-package: 'jre'
+      - name: Fetch Wiremock
+        shell: bash
+        run: curl https://repo1.maven.org/maven2/org/wiremock/wiremock-standalone/3.11.0/wiremock-standalone-3.11.0.jar --output .wiremock/wiremock-standalone.jar
       - name: Setup parameters file
         shell: bash
         env:

DESCRIPTION.md

@@ -15,6 +15,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Fixed a bug where privatelink OCSP Cache url could not be determined if privatelink account name was specified in uppercase.
   - Added support for iceberg tables to `write_pandas`.
   - Fixed base64 encoded private key tests.
+  - Added Wiremock tests.
 
 - v3.13.2(January 29, 2025)
   - Changed not to use scoped temporary objects.

MANIFEST.in

@@ -19,6 +19,7 @@ exclude license_header.txt
 exclude tox.ini
 exclude mypy.ini
 exclude .clang-format
+exclude .wiremock/*
 
 prune ci
 prune benchmark

ci/docker/connector_test_fips/Dockerfile

@@ -19,6 +19,7 @@ RUN sed -i s/mirror.centos.org/vault.centos.org/g /etc/yum.repos.d/*.repo && \
 RUN yum clean all && \
     yum install -y redhat-rpm-config gcc libffi-devel openssl openssl-devel && \
     yum install -y python38 python38-devel && \
+    yum install -y java-11-openjdk && \
     yum clean all && \
     rm -rf /var/cache/yum
 RUN python3 -m pip install --user --upgrade pip setuptools wheel

ci/test_fips.sh

@@ -8,6 +8,9 @@ THIS_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
 CONNECTOR_DIR="$( dirname "${THIS_DIR}")"
 CONNECTOR_WHL="$(ls $CONNECTOR_DIR/dist/*cp38*manylinux2014*.whl | sort -r | head -n 1)"
 
+# fetch wiremock
+curl https://repo1.maven.org/maven2/org/wiremock/wiremock-standalone/3.11.0/wiremock-standalone-3.11.0.jar --output "${CONNECTOR_DIR}/.wiremock/wiremock-standalone.jar"
+
 python3.8 -m venv fips_env
 source fips_env/bin/activate
 pip install -U setuptools pip

test/unit/test_wiremock_client.py

@@ -0,0 +1,52 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Any, Generator, Union
+
+import pytest
+
+# old driver support
+try:
+    from snowflake.connector.vendored import requests
+    from src.snowflake.connector.test_util import RUNNING_ON_JENKINS
+except ImportError:
+    import os
+
+    import requests
+
+    RUNNING_ON_JENKINS = os.getenv("JENKINS_HOME") is not None
+
+
+from ..wiremock.wiremock_utils import WiremockClient
+
+
+@pytest.fixture(scope="session")
+def wiremock_client() -> Generator[Union[WiremockClient, Any], Any, None]:
+    with WiremockClient() as client:
+        yield client
+
+
+@pytest.mark.skipif(RUNNING_ON_JENKINS, reason="jenkins doesn't support wiremock tests")
+def test_wiremock(wiremock_client):
+    connection_reset_by_peer_mapping = {
+        "mappings": [
+            {
+                "scenarioName": "Basic example",
+                "requiredScenarioState": "Started",
+                "request": {"method": "GET", "url": "/endpoint"},
+                "response": {"status": 200},
+            }
+        ],
+        "importOptions": {"duplicatePolicy": "IGNORE", "deleteAllNotInImport": True},
+    }
+    wiremock_client.import_mapping(connection_reset_by_peer_mapping)
+
+    response = requests.get(
+        f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/endpoint"
+    )
+
+    assert response is not None, "response is None"
+    assert (
+        response.status_code == requests.codes.ok
+    ), f"response status is not 200, received status {response.status_code}"

test/wiremock/__init__.py

@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#

test/wiremock/wiremock_utils.py

@@ -0,0 +1,152 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+import json
+import logging
+import os.path
+import pathlib
+import socket
+import subprocess
+from time import sleep
+from typing import Optional
+
+try:
+    from snowflake.connector.vendored import requests
+except ImportError:
+    import requests
+
+WIREMOCK_START_MAX_RETRY_COUNT = 12
+LOGGER = logging.getLogger(__name__)
+
+
+def _get_mapping_str(mapping):
+    if isinstance(mapping, str):
+        return mapping
+    if isinstance(mapping, dict):
+        return json.dumps(mapping)
+    if os.path.isfile(str(mapping)):
+        with open(mapping) as f:
+            return f.read()
+
+    raise RuntimeError(f"Mapping {mapping} is of an invalid type")
+
+
+class WiremockClient:
+    def __init__(self):
+        self.wiremock_filename = "wiremock-standalone.jar"
+        self.wiremock_host = "localhost"
+        self.wiremock_http_port = None
+        self.wiremock_https_port = None
+
+        self.wiremock_dir = pathlib.Path(__file__).parent.parent.parent / ".wiremock"
+        assert self.wiremock_dir.exists(), f"{self.wiremock_dir} does not exist"
+
+        self.wiremock_jar_path = self.wiremock_dir / self.wiremock_filename
+        assert (
+            self.wiremock_jar_path.exists()
+        ), f"{self.wiremock_jar_path} does not exist"
+
+    def _start_wiremock(self):
+        self.wiremock_http_port = self._find_free_port()
+        self.wiremock_https_port = self._find_free_port()
+        self.wiremock_process = subprocess.Popen(
+            [
+                "java",
+                "-jar",
+                self.wiremock_jar_path,
+                "--root-dir",
+                self.wiremock_dir,
+                "--enable-browser-proxying",  # work as forward proxy
+                "--proxy-pass-through",
+                "false",  # pass through only matched requests
+                "--port",
+                str(self.wiremock_http_port),
+                "--https-port",
+                str(self.wiremock_https_port),
+                "--https-keystore",
+                self.wiremock_dir / "ca-cert.jks",
+                "--ca-keystore",
+                self.wiremock_dir / "ca-cert.jks",
+            ]
+        )
+        self._wait_for_wiremock()
+
+    def _stop_wiremock(self):
+        self.wiremock_process.kill()
+
+    def _wait_for_wiremock(self):
+        retry_count = 0
+        while retry_count < WIREMOCK_START_MAX_RETRY_COUNT:
+            if self._health_check():
+                return
+            retry_count += 1
+            sleep(1)
+
+        raise TimeoutError(
+            f"WiremockClient did not respond within {WIREMOCK_START_MAX_RETRY_COUNT} seconds"
+        )
+
+    def _health_check(self):
+        mappings_endpoint = (
+            f"http://{self.wiremock_host}:{self.wiremock_http_port}/__admin/health"
+        )
+        try:
+            response = requests.get(mappings_endpoint)
+        except requests.exceptions.RequestException as e:
+            LOGGER.warning(f"Wiremock healthcheck failed with exception: {e}")
+            return False
+        if (
+            response.status_code == requests.codes.ok
+            and response.json()["status"] == "healthy"
+        ):
+            LOGGER.debug(f"Wiremock healthcheck failed with response: {response}")
+        else:
+            LOGGER.warning(
+                f"Wiremock healthcheck failed with status code: {response.status_code}"
+            )
+            return False
+        return True
+
+    def _reset_wiremock(self):
+        reset_endpoint = (
+            f"http://{self.wiremock_host}:{self.wiremock_http_port}/__admin/reset"
+        )
+        response = self._wiremock_post(reset_endpoint)
+        if response.status_code != requests.codes.ok:
+            raise RuntimeError("Failed to reset WiremockClient")
+
+    def _wiremock_post(
+        self, endpoint: str, body: Optional[str] = None
+    ) -> requests.Response:
+        headers = {"Accept": "application/json", "Content-Type": "application/json"}
+        return requests.post(endpoint, data=body, headers=headers)
+
+    def import_mapping(self, mapping):
+        self._reset_wiremock()
+        import_mapping_endpoint = f"http://{self.wiremock_host}:{self.wiremock_http_port}/__admin/mappings/import"
+        mapping_str = _get_mapping_str(mapping)
+        response = self._wiremock_post(import_mapping_endpoint, mapping_str)
+        if response.status_code != requests.codes.ok:
+            raise RuntimeError("Failed to import mapping")
+
+    def add_mapping(self, mapping):
+        add_mapping_endpoint = (
+            f"http://{self.wiremock_host}:{self.wiremock_http_port}/__admin/mappings"
+        )
+        mapping_str = _get_mapping_str(mapping)
+        response = self._wiremock_post(add_mapping_endpoint, mapping_str)
+        if response.status_code != requests.codes.created:
+            raise RuntimeError("Failed to add mapping")
+
+    def _find_free_port(self) -> int:
+        with socket.socket() as sock:
+            sock.bind((self.wiremock_host, 0))
+            return sock.getsockname()[1]
+
+    def __enter__(self):
+        self._start_wiremock()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._stop_wiremock()