APP-6435: Added support for vcrpy test utilities to mock HTTP interactions with 3rd-party APIs

Aryamanz29 · Aryamanz29 · commit effd6c859142 · 2025-04-28T16:32:50.000+05:30
diff --git a/pyatlan/test_utils/__init__.py b/pyatlan/test_utils/__init__.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2024 Atlan Pte. Ltd.
+# Copyright 2025 Atlan Pte. Ltd.
 import logging
 import random
 from os import path
diff --git a/pyatlan/test_utils/base_vcr.py b/pyatlan/test_utils/base_vcr.py
@@ -0,0 +1,280 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+
+import pkg_resources  # type: ignore[import-untyped]
+
+from pyatlan.utils import DependencyNotFoundError
+
+try:
+    # Check if pytest-vcr plugin is installed
+    pkg_resources.get_distribution("pytest-vcr")
+except pkg_resources.DistributionNotFound:
+    raise DependencyNotFoundError("pytest-vcr")
+
+import json
+import os
+from typing import Any, Dict, Union
+
+import pytest
+import yaml  # type: ignore[import-untyped]
+
+
+class LiteralBlockScalar(str):
+    """Formats the string as a literal block scalar, preserving whitespace and
+    without interpreting escape characters"""
+
+
+def literal_block_scalar_presenter(dumper, data):
+    """Represents a scalar string as a literal block, via '|' syntax"""
+    return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
+
+
+yaml.add_representer(LiteralBlockScalar, literal_block_scalar_presenter)
+
+
+def process_string_value(string_value):
+    """Pretty-prints JSON or returns long strings as a LiteralBlockScalar"""
+    try:
+        json_data = json.loads(string_value)
+        return LiteralBlockScalar(json.dumps(json_data, indent=2))
+    except (ValueError, TypeError):
+        if len(string_value) > 80:
+            return LiteralBlockScalar(string_value)
+    return string_value
+
+
+def convert_body_to_literal(data):
+    """Searches the data for body strings, attempting to pretty-print JSON"""
+    if isinstance(data, dict):
+        for key, value in data.items():
+            # Handle response body case (e.g: response.body.string)
+            if key == "body" and isinstance(value, dict) and "string" in value:
+                value["string"] = process_string_value(value["string"])
+
+            # Handle request body case (e.g: request.body)
+            elif key == "body" and isinstance(value, str):
+                data[key] = process_string_value(value)
+
+            else:
+                convert_body_to_literal(value)
+
+    elif isinstance(data, list):
+        for idx, choice in enumerate(data):
+            data[idx] = convert_body_to_literal(choice)
+
+    return data
+
+
+class VCRPrettyPrintYamlJSONBody:
+    """This makes request and response YAML JSON body recordings more readable."""
+
+    @staticmethod
+    def serialize(cassette_dict):
+        cassette_dict = convert_body_to_literal(cassette_dict)
+        return yaml.dump(cassette_dict, default_flow_style=False, allow_unicode=True)
+
+    @staticmethod
+    def deserialize(cassette_string):
+        return yaml.safe_load(cassette_string)
+
+
+class VCRPrettyPrintJSONBody:
+    """Makes request and response JSON body recordings more readable."""
+
+    @staticmethod
+    def _parse_json_body(
+        body: Union[str, bytes, None],
+    ) -> Union[Dict[str, Any], str, None, bytes]:
+        """Parse JSON body if possible, otherwise return the original body."""
+        if body is None:
+            return None
+
+        # Convert bytes to string if needed
+        if isinstance(body, bytes):
+            try:
+                body = body.decode("utf-8")
+            except UnicodeDecodeError:
+                return body  # Return original if can't decode
+
+        # If it's a string, try to parse as JSON
+        if isinstance(body, str):
+            try:
+                return json.loads(body)
+            except json.JSONDecodeError:
+                return body  # Return original if not valid JSON
+
+        return body  # Return original for other types
+
+    @staticmethod
+    def serialize(cassette_dict: dict) -> str:
+        """
+        Converts body strings to parsed JSON objects for better readability when possible.
+        """
+        # Safety check for cassette_dict
+        if not cassette_dict or not isinstance(cassette_dict, dict):
+            cassette_dict = {}
+
+        interactions = cassette_dict.get("interactions", []) or []
+
+        for interaction in interactions:
+            if not interaction:
+                continue
+
+            # Handle response body
+            response = interaction.get("response") or {}
+            body_container = response.get("body")
+            if isinstance(body_container, dict) and "string" in body_container:
+                parsed_body = VCRPrettyPrintJSONBody._parse_json_body(
+                    body_container["string"]
+                )
+                if isinstance(parsed_body, dict):
+                    # Replace string field with parsed_json field
+                    response["body"] = {"parsed_json": parsed_body}
+
+            # Handle request body
+            request = interaction.get("request") or {}
+            body_container = request.get("body")
+            if isinstance(body_container, dict) and "string" in body_container:
+                parsed_body = VCRPrettyPrintJSONBody._parse_json_body(
+                    body_container["string"]
+                )
+                if isinstance(parsed_body, dict):
+                    # Replace string field with parsed_json field
+                    request["body"] = {"parsed_json": parsed_body}
+
+        # Serialize the final dictionary into a JSON string with pretty formatting
+        try:
+            return json.dumps(cassette_dict, indent=2, ensure_ascii=False) + "\n"
+        except TypeError as exc:
+            raise TypeError(
+                "Does this HTTP interaction contain binary data? "
+                "If so, use a different serializer (like the YAML serializer)."
+            ) from exc
+
+    @staticmethod
+    def deserialize(cassette_string: str) -> dict:
+        """
+        Deserializes a JSON string into a dictionary and converts
+        parsed_json fields back to string fields.
+        """
+        # Safety check for cassette_string
+        if not cassette_string:
+            return {}
+
+        try:
+            cassette_dict = json.loads(cassette_string)
+        except json.JSONDecodeError:
+            return {}
+
+        # Convert parsed_json back to string format
+        interactions = cassette_dict.get("interactions", []) or []
+
+        for interaction in interactions:
+            if not interaction:
+                continue
+
+            # Handle response body
+            response = interaction.get("response") or {}
+            body_container = response.get("body")
+            if isinstance(body_container, dict) and "parsed_json" in body_container:
+                json_body = body_container["parsed_json"]
+                response["body"] = {"string": json.dumps(json_body)}
+
+            # Handle request body
+            request = interaction.get("request") or {}
+            body_container = request.get("body")
+            if isinstance(body_container, dict) and "parsed_json" in body_container:
+                json_body = body_container["parsed_json"]
+                request["body"] = {"string": json.dumps(json_body)}
+
+        return cassette_dict
+
+
+class VCRRemoveAllHeaders:
+    """
+    A class responsible for removing all headers from requests and responses.
+    This can be useful for scenarios where headers are not needed for matching or comparison
+    in VCR (Virtual Cassette Recorder) interactions, such as when recording or replaying HTTP requests.
+    """
+
+    @staticmethod
+    def remove_all_request_headers(request):
+        # Save only what's necessary for matching
+        request.headers = {}
+        return request
+
+    @staticmethod
+    def remove_all_response_headers(response):
+        # Save only what's necessary for matching
+        response["headers"] = {}
+        return response
+
+
+class BaseVCR:
+    """
+    A base class for configuring VCR (Virtual Cassette Recorder)
+    for HTTP request/response recording and replaying.
+
+    This class provides pytest fixtures to set up the VCR configuration
+    and custom serializers for JSON and YAML formats.
+    It also handles cassette directory configuration.
+    """
+
+    _CASSETTES_DIR = None
+
+    @pytest.fixture(scope="module")
+    def vcr(self, vcr):
+        """
+        Registers custom serializers for VCR and returns the VCR instance.
+
+        The method registers two custom serializers:
+        - "pretty-json" for pretty-printing JSON responses.
+        - "pretty-yaml" for pretty-printing YAML responses.
+
+        :param vcr: The VCR instance provided by the pytest-vcr plugin
+        :returns: modified VCR instance with custom serializers registered
+        """
+        vcr.register_serializer("pretty-json", VCRPrettyPrintJSONBody)
+        vcr.register_serializer("pretty-yaml", VCRPrettyPrintYamlJSONBody)
+        return vcr
+
+    @pytest.fixture(scope="module")
+    def vcr_config(self):
+        """
+        Provides the VCR configuration dictionary.
+
+        The configuration includes default options for the recording mode,
+        serializer, response decoding, and filtering headers.
+        This configuration is used to set up VCR behavior during tests.
+
+        :returns: a dictionary with VCR configuration options
+        """
+        return {
+            # More config options can be found at:
+            # https://vcrpy.readthedocs.io/en/latest/configuration.html#configuration
+            "record_mode": "once",  # (default: "once", "always", "none", "new_episodes")
+            "serializer": "pretty-yaml",  # (default: "yaml")
+            "decode_compressed_response": True,  # Decode compressed responses
+            # (optional) Replace the Authorization request header with "**REDACTED**" in cassettes
+            # "filter_headers": [("authorization", "**REDACTED**")],
+            "before_record_request": VCRRemoveAllHeaders.remove_all_request_headers,
+            "before_record_response": VCRRemoveAllHeaders.remove_all_response_headers,
+        }
+
+    @pytest.fixture(scope="module")
+    def vcr_cassette_dir(self, request):
+        """
+        Provides the directory path for storing VCR cassettes.
+
+        If a custom cassette directory is set in the class, it is used;
+        otherwise, the default directory structure is created under "tests/cassettes".
+        The directory path will be based on the module name.
+
+        :param request: request object which provides metadata about the test
+
+        :returns: directory path for storing cassettes
+        """
+        # Set self._CASSETTES_DIR or use the default directory path based on the test module name
+        return self._CASSETTES_DIR or os.path.join(
+            "tests/cassettes", request.module.__name__
+        )
diff --git a/pyatlan/utils.py b/pyatlan/utils.py
@@ -471,3 +471,10 @@ def validate_single_required_field(field_names: List[str], values: List[Any]):
         raise ValueError(
             f"Only one of the following parameters are allowed: {', '.join(names)}"
         )
+
+
+class DependencyNotFoundError(Exception):
+    def __init__(self, dependency):
+        super().__init__(
+            f"{dependency} is not installed, but it is required to use this module. Please install {dependency}."
+        )
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -2,6 +2,7 @@ mypy~=1.9.0
 ruff~=0.9.9
 types-requests~=2.32.0.20241016
 pytest~=8.3.4
+pytest-vcr~=1.0.2
 pytest-order~=1.3.0
 pytest-timer[termcolor]~=1.0.0
 pytest-sugar~=1.0.0
