ddapm-test-agent 1.37.0__py3-none-any.whl → 1.38.0__py3-none-any.whl

ddapm_test_agent/vcr_proxy.py

@@ -1,4 +1,7 @@
 import asyncio
+import base64
+from dataclasses import dataclass
+from glob import glob
 import hashlib
 import json
 import logging
@@ -6,29 +9,105 @@ import os
 import re
 from typing import Any
 from typing import Dict
+from typing import List
 from typing import Optional
+from typing import Union
+from typing import cast
 from urllib.parse import urljoin
 
 from aiohttp.web import Request
 from aiohttp.web import Response
 import requests
 from requests_aws4auth import AWS4Auth
-import vcr
+import yaml
 
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class CassetteDataRequest:
+    """Represents the request portion of a cassette."""
+
+    method: str
+    url: str
+    headers: Dict[str, str]
+    body: str
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CassetteDataRequest":
+        """Create from a dictionary."""
+        return cls(
+            method=data["method"],
+            url=data["url"],
+            headers=data["headers"],
+            body=data["body"],
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to a dictionary for JSON serialization."""
+        return {
+            "method": self.method,
+            "url": self.url,
+            "headers": self.headers,
+            "body": self.body,
+        }
+
+
+@dataclass
+class CassetteDataResponse:
+    """Represents the response portion of a cassette."""
+
+    status: Dict[str, Any]  # {"code": int, "message": str}
+    headers: Dict[str, str]
+    body: str
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CassetteDataResponse":
+        """Create from a dictionary."""
+        return cls(
+            status=data["status"],
+            headers=data["headers"],
+            body=data["body"],
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to a dictionary for JSON serialization."""
+        return {
+            "status": self.status,
+            "headers": self.headers,
+            "body": self.body,
+        }
+
+
+@dataclass
+class CassetteData:
+    """Represents a VCR cassette with request and response data."""
+
+    request: CassetteDataRequest
+    response: CassetteDataResponse
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CassetteData":
+        """Create from a dictionary (e.g., loaded from JSON)."""
+        return cls(
+            request=CassetteDataRequest.from_dict(data["request"]),
+            response=CassetteDataResponse.from_dict(data["response"]),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to a dictionary for JSON serialization."""
+        return {
+            "request": self.request.to_dict(),
+            "response": self.response.to_dict(),
+        }
+
+
 #  Used for AWS signature recalculation for aws services initial proxying
 AWS_REGION = os.environ.get("AWS_REGION", "us-east-1")
 AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY")
 
 
-def url_path_join(base_url: str, path: str) -> str:
-    """Join a base URL with a path, handling slashes automatically."""
-    return urljoin(base_url.rstrip("/") + "/", path.lstrip("/"))
-
-
 AWS_SERVICES = {
     "bedrock-runtime": "bedrock",
 }
@@ -72,11 +151,16 @@ NORMALIZERS = [
 ]
 
 
+def _url_path_join(base_url: str, path: str) -> str:
+    """Join a base URL with a path, handling slashes automatically."""
+    return urljoin(base_url.rstrip("/") + "/", path.lstrip("/"))
+
+
 def _file_safe_string(s: str) -> str:
     return "".join(c if c.isalnum() or c in ".-" else "_" for c in s)
 
 
-def get_custom_vcr_providers(vcr_provider_map: str) -> Dict[str, str]:
+def _get_custom_vcr_providers(vcr_provider_map: str) -> Dict[str, str]:
     return dict(
         [
             vcr_provider_map.strip().split("=", 1)
@@ -86,7 +170,7 @@ def get_custom_vcr_providers(vcr_provider_map: str) -> Dict[str, str]:
     )
 
 
-def normalize_multipart_body(body: bytes) -> str:
+def _normalize_multipart_body(body: bytes) -> str:
     if not body:
         return ""
 
@@ -110,7 +194,41 @@ def normalize_multipart_body(body: bytes) -> str:
             return f"[binary_data_{hex_digest}]"
 
 
-def parse_authorization_header(auth_header: str) -> Dict[str, str]:
+def _decode_body(body: bytes) -> str:
+    """Decode body (request or response), handling binary data gracefully."""
+    if not body:
+        return ""
+
+    # Check for null bytes - strong indicator of binary data (e.g., event streams, protobuf)
+    if b"\x00" in body:
+        return "base64:" + base64.b64encode(body).decode("ascii")
+
+    try:
+        # Try UTF-8 decode - if successful, it's text
+        return body.decode("utf-8")
+    except UnicodeDecodeError:
+        # If UTF-8 fails, treat as binary
+        return "base64:" + base64.b64encode(body).decode("ascii")
+
+
+def _encode_body(body: str) -> bytes:
+    """Convert cassette body string back to bytes, handling base64-encoded data."""
+    if not body:
+        return b""
+
+    # Check for base64 marker first (for binary data that was base64-encoded)
+    if body.startswith("base64:"):
+        return base64.b64decode(body[7:])
+
+    try:
+        # Try to encode as UTF-8 (most common case)
+        return body.encode("utf-8")
+    except UnicodeEncodeError:
+        # If all else fails, encode as latin-1
+        return body.encode("latin-1")
+
+
+def _parse_authorization_header(auth_header: str) -> Dict[str, str]:
     """Parse AWS Authorization header to extract components"""
     if not auth_header.startswith("AWS4-HMAC-SHA256 "):
         return {}
@@ -125,20 +243,8 @@ def parse_authorization_header(auth_header: str) -> Dict[str, str]:
     return parsed
 
 
-def get_vcr(subdirectory: str, vcr_cassettes_directory: str, vcr_ignore_headers: str) -> vcr.VCR:
-    cassette_dir = os.path.join(vcr_cassettes_directory, subdirectory)
-    extra_ignore_headers = vcr_ignore_headers.split(",")
-
-    return vcr.VCR(
-        cassette_library_dir=cassette_dir,
-        record_mode="once",
-        match_on=["path", "method"],
-        filter_headers=CASSETTE_FILTER_HEADERS + extra_ignore_headers,
-    )
-
-
-def generate_cassette_name(path: str, method: str, body: bytes, vcr_cassette_prefix: Optional[str]) -> str:
-    decoded_body = normalize_multipart_body(body) if body else ""
+def _generate_cassette_name(path: str, method: str, body: bytes, vcr_cassette_prefix: Optional[str]) -> str:
+    decoded_body = _normalize_multipart_body(body) if body else ""
     try:
         parsed_body = json.loads(decoded_body) if decoded_body else {}
     except json.JSONDecodeError:
@@ -158,11 +264,180 @@ def generate_cassette_name(path: str, method: str, body: bytes, vcr_cassette_pre
     )
 
 
+def _convert_vcr_cassette_to_custom_format(
+    cassette_file_path: str,
+    request_kwargs: Dict[str, Any],
+    vcr_ignore_headers: str,
+) -> CassetteData:
+    """Convert a VCR YAML cassette to our custom JSON format."""
+    cassette_file_path_yaml = f"{cassette_file_path}.yaml"
+    with open(cassette_file_path_yaml, "r") as f:
+        cassette_data = yaml.load(f, Loader=yaml.UnsafeLoader)
+
+    interaction = cast(Dict[str, Any], cassette_data["interactions"][0])
+
+    cassette = _write_cassette_file(cassette_file_path, request_kwargs, interaction["response"], vcr_ignore_headers)
+
+    logger.warning(f"Removing legacy VCR cassette file {cassette_file_path_yaml}.")
+    os.remove(cassette_file_path_yaml)
+
+    return cassette
+
+
+def _normalize_header_value(value: Any) -> str:
+    """Normalize header value to a string (handles list values)."""
+    if isinstance(value, list):
+        return str(value[0]) if value else ""
+    return str(value)
+
+
+def _filter_headers(headers: Dict[str, Any], ignore_headers: List[str]) -> Dict[str, str]:
+    """Filter headers and normalize their values."""
+    return {key: _normalize_header_value(value) for key, value in headers.items() if key.lower() not in ignore_headers}
+
+
+def _create_cassette_from_requests_response(
+    request_kwargs: Dict[str, Any],
+    response: requests.Response,
+    ignore_headers: List[str],
+) -> CassetteData:
+    """Create cassette data from a requests.Response object."""
+    logger.info(f"Creating cassette data from requests.Response object: {response.content!r}")
+    return CassetteData(
+        request=CassetteDataRequest(
+            method=request_kwargs["method"],
+            url=request_kwargs["url"],
+            headers=_filter_headers(request_kwargs["headers"], ignore_headers),
+            body=_decode_body(request_kwargs["data"]) if request_kwargs["data"] else "",
+        ),
+        response=CassetteDataResponse(
+            status={
+                "code": response.status_code,
+                "message": response.reason or "",
+            },
+            headers=_filter_headers(dict(response.headers), ignore_headers),
+            body=_decode_body(response.content) if response.content else "",
+        ),
+    )
+
+
+def _create_cassette_from_dict(
+    request_kwargs: Dict[str, Any],
+    response_dict: Dict[str, Any],
+    ignore_headers: List[str],
+) -> CassetteData:
+    """Create cassette data from a dictionary (e.g., from VCR YAML)."""
+    body_data = response_dict["body"]["string"]
+    if isinstance(body_data, bytes):
+        body_str = _decode_body(body_data)
+    else:
+        body_str = body_data
+
+    return CassetteData(
+        request=CassetteDataRequest(
+            method=request_kwargs["method"],
+            url=request_kwargs["url"],
+            headers=_filter_headers(request_kwargs["headers"], ignore_headers),
+            body=_decode_body(request_kwargs["data"]) if request_kwargs["data"] else "",
+        ),
+        response=CassetteDataResponse(
+            status={
+                "code": response_dict["status"]["code"],
+                "message": response_dict["status"]["message"],
+            },
+            headers=_filter_headers(response_dict["headers"], ignore_headers),
+            body=body_str,
+        ),
+    )
+
+
+def _write_cassette_file(
+    cassette_file_path: str,
+    request_kwargs: Dict[str, Any],
+    response: Union[requests.Response, Dict[str, Any]],
+    vcr_ignore_headers: str,
+) -> CassetteData:
+    """Write cassette data to a JSON file."""
+    cassette_file_path_json = f"{cassette_file_path}.json"
+    logger.info(f"Writing cassette file to {cassette_file_path_json}")
+
+    cassette_dir = os.path.dirname(cassette_file_path_json)
+    os.makedirs(cassette_dir, exist_ok=True)
+
+    ignore_headers_list = [
+        header.lower() for header in CASSETTE_FILTER_HEADERS + vcr_ignore_headers.split(",") if header
+    ]
+
+    if isinstance(response, requests.Response):
+        cassette = _create_cassette_from_requests_response(request_kwargs, response, ignore_headers_list)
+    else:
+        # conversion of legacy VCR cassette to JSON format
+        cassette = _create_cassette_from_dict(request_kwargs, response, ignore_headers_list)
+
+    with open(cassette_file_path_json, "w") as f:
+        json.dump(cassette.to_dict(), f, indent=2)
+
+    return cassette
+
+
+async def _request(
+    cassette_file_path: str, cassette_exists: bool, request_kwargs: Dict[str, Any], vcr_ignore_headers: str
+) -> Response:
+    """
+    Load a cassette from file if it exists, otherwise make a request and save the response.
+
+    If the cassette was created with the VCR package (YAML format), convert it to JSON format.
+    """
+    logger.info(f"Making a request to {request_kwargs['url']} with method {request_kwargs['method']}")
+
+    cassette: Optional[CassetteData] = None
+
+    if cassette_exists:
+        logger.info(f"Cassette file exists at {cassette_file_path}")
+        cassette_files = glob(f"{cassette_file_path}.*")
+        if not cassette_files:
+            raise FileNotFoundError(f"Expected cassette file at {cassette_file_path}.*")
+
+        file_extension = os.path.splitext(cassette_files[0])[1]
+
+        if file_extension == ".yaml":  # TODO(sabrenner): in a breaking change, remove this
+            logger.warning(
+                "Converting legacy VCR cassette to JSON format. This will not be supported in ddapm-test-agent==2.0.0"
+            )
+            cassette = _convert_vcr_cassette_to_custom_format(cassette_file_path, request_kwargs, vcr_ignore_headers)
+        elif file_extension == ".json":
+            cassette_file_path_json = f"{cassette_file_path}.json"
+            with open(cassette_file_path_json, "r") as f:
+                cassette = CassetteData.from_dict(json.load(f))
+        else:
+            raise ValueError(f"Unsupported cassette file extension: {file_extension}")
+    else:
+        logger.info(f"Cassette file does not exist at {cassette_file_path}, making a request to the provider")
+        provider_response = await asyncio.to_thread(lambda: requests.request(**request_kwargs))
+        cassette = _write_cassette_file(cassette_file_path, request_kwargs, provider_response, vcr_ignore_headers)
+
+    # Build response from cassette data
+    response_body_str = cassette.response.body
+    response_body = _encode_body(response_body_str) if isinstance(response_body_str, str) else b""
+
+    response = Response(
+        body=response_body,
+        status=cassette.response.status["code"],
+    )
+
+    skip_headers = {"content-length", "transfer-encoding", "content-encoding", "connection"}
+    for key, value in cassette.response.headers.items():
+        if key.lower() not in skip_headers:
+            response.headers[key] = value
+
+    return response
+
+
 async def proxy_request(
     request: Request, vcr_cassettes_directory: str, vcr_ci_mode: bool, vcr_provider_map: str, vcr_ignore_headers: str
 ) -> Response:
     provider_base_urls = PROVIDER_BASE_URLS.copy()
-    provider_base_urls.update(get_custom_vcr_providers(vcr_provider_map))
+    provider_base_urls.update(_get_custom_vcr_providers(vcr_provider_map))
 
     path = request.match_info["path"]
     if request.query_string:
@@ -179,18 +454,17 @@ async def proxy_request(
     body_bytes = await request.read()
 
     vcr_cassette_prefix = request.pop("vcr_cassette_prefix", None)
-    cassette_name = generate_cassette_name(path, request.method, body_bytes, vcr_cassette_prefix)
-    cassette_file_name = f"{cassette_name}.yaml"
-    cassette_file_path = os.path.join(vcr_cassettes_directory, provider, cassette_file_name)
-    cassette_exists = os.path.exists(cassette_file_path)
+    cassette_name = _generate_cassette_name(path, request.method, body_bytes, vcr_cassette_prefix)
+    cassette_file_path = os.path.join(vcr_cassettes_directory, provider, cassette_name)
+    cassette_exists = len(glob(f"{cassette_file_path}.*")) > 0
 
     if vcr_ci_mode and not cassette_exists:
         return Response(
-            body=f"Cassette {cassette_file_name} not found while running in CI mode. Please generate the cassette locally and commit it.",
+            body=f"Cassette {cassette_name} not found while running in CI mode. Please generate the cassette locally and commit it.",
             status=500,
         )
 
-    target_url = url_path_join(provider_base_urls[provider], remaining_path)
+    target_url = _url_path_join(provider_base_urls[provider], remaining_path)
     headers = {key: value for key, value in request.headers.items() if key != "Host"}
 
     request_kwargs: Dict[str, Any] = {
@@ -211,36 +485,10 @@ async def proxy_request(
             )
 
         auth_header = request.headers.get("Authorization", "")
-        auth_parts = parse_authorization_header(auth_header)
+        auth_parts = _parse_authorization_header(auth_header)
         aws_access_key = auth_parts.get("Credential", "").split("/")[0]
 
         auth = AWS4Auth(aws_access_key, AWS_SECRET_ACCESS_KEY, AWS_REGION, AWS_SERVICES[provider])
         request_kwargs["auth"] = auth
 
-    def _make_request():
-        with get_vcr(provider, vcr_cassettes_directory, vcr_ignore_headers).use_cassette(cassette_file_name):
-            return requests.request(**request_kwargs)
-
-    provider_response = await asyncio.to_thread(_make_request)
-
-    # Extract content type without charset
-    content_type = provider_response.headers.get("content-type", "")
-    if ";" in content_type:
-        content_type = content_type.split(";")[0].strip()
-
-    response = Response(
-        body=provider_response.content,
-        status=provider_response.status_code,
-        content_type=content_type,
-    )
-
-    for key, value in provider_response.headers.items():
-        if key.lower() not in (
-            "content-length",
-            "transfer-encoding",
-            "content-encoding",
-            "connection",
-        ):
-            response.headers[key] = value
-
-    return response
+    return await _request(cassette_file_path, cassette_exists, request_kwargs, vcr_ignore_headers)

ddapm_test_agent/web.py

@@ -797,7 +797,7 @@ class WebUI:
 
             # Get flares for all other session tokens
             session_tokens = set()
-            for req_data in (request_storage.get_all_requests() if request_storage is not None else []):
+            for req_data in request_storage.get_all_requests() if request_storage is not None else []:
                 token = req_data.get("headers", {}).get("X-Datadog-Test-Session-Token")
                 if token:
                     session_tokens.add(token)
@@ -934,7 +934,7 @@ class WebUI:
 
             # Get flares for all other session tokens
             session_tokens = set()
-            for req_data in (request_storage.get_all_requests() if request_storage is not None else []):
+            for req_data in request_storage.get_all_requests() if request_storage is not None else []:
                 token = req_data.get("headers", {}).get("X-Datadog-Test-Session-Token")
                 if token:
                     session_tokens.add(token)

{ddapm_test_agent-1.37.0.dist-info → ddapm_test_agent-1.38.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ddapm-test-agent
-Version: 1.37.0
+Version: 1.38.0
 Summary: Test agent for Datadog APM client libraries
 Home-page: https://github.com/Datadog/dd-apm-test-agent
 Author: Kyle Verhoog
@@ -20,9 +20,9 @@ Requires-Dist: msgpack
 Requires-Dist: requests
 Requires-Dist: typing_extensions
 Requires-Dist: yarl
-Requires-Dist: vcrpy
 Requires-Dist: requests-aws4auth
 Requires-Dist: jinja2>=3.0.0
+Requires-Dist: pyyaml
 Requires-Dist: opentelemetry-proto<1.37.0,>1.33.0
 Requires-Dist: protobuf>=3.19.0
 Requires-Dist: grpcio<2.0,>=1.66.2

{ddapm_test_agent-1.37.0.dist-info → ddapm_test_agent-1.38.0.dist-info}/RECORD

@@ -16,8 +16,8 @@ ddapm_test_agent/trace_snapshot.py,sha256=vcz9uCgtpnInKl32nq1n62shhsVdMQPzOWfV3-
 ddapm_test_agent/tracerflare.py,sha256=uoSjhPCOKZflgJn5JLv1Unh4gUdAR1-YbC9_1n1iH9w,954
 ddapm_test_agent/tracestats.py,sha256=q_WQZnh2kXSSN3fRIBe_0jMYCBQHcaS3fZmJTge4lWc,2073
 ddapm_test_agent/tracestats_snapshot.py,sha256=VsB6MVnHPjPWHVWnnDdCXJcVKL_izKXEf9lvJ0qbjNQ,3609
-ddapm_test_agent/vcr_proxy.py,sha256=dEJ5xXxxvFHAbvN-OgXEaOlqz54Vh1k_1eQBgf146lU,7812
-ddapm_test_agent/web.py,sha256=cvQqncgLbciKpG7lyLJRuNbKA9rp_mhuRXWFwuom1L0,59768
+ddapm_test_agent/vcr_proxy.py,sha256=zbenLsLVBStVJlBFypf0fFuZcDLMpXDT-GJ1ts-tCE0,16654
+ddapm_test_agent/web.py,sha256=D57xkEHuIMSi_ZVEIpxlYEbt_6cipAVhz1XneVWd8AU,59764
 ddapm_test_agent/static/style.css,sha256=0sUpN62bhc3cuW89G39hzNBmIPaYc8MPBt_PvSn0zf0,29560
 ddapm_test_agent/templates/base.html,sha256=Ie9ZvQXtv7YSR4SMFYUp5-eGs5sB_fxApU5pl1jUCCc,1010
 ddapm_test_agent/templates/config.html,sha256=vIgIG79Wo95p5x8RL0Esr9SZ86Ueejt9OBmW_uev1-o,12440
@@ -31,10 +31,10 @@ ddapm_test_agent/templates/snapshots.html,sha256=uWiYK9yDP9S88nxYNr0vLhYkWSzzj9C
 ddapm_test_agent/templates/trace_detail.html,sha256=EqC6l_jo2Afrf511_-pEDgz6-TEogTjtjawMVxzw7G4,1235
 ddapm_test_agent/templates/tracer_flares.html,sha256=wwq7Yaq7shXJineHT38A2tIHIEoirs5ryWu7raFco80,16272
 ddapm_test_agent/templates/traces.html,sha256=-RB_Nc3NzQQXApw5Y37KbmkU4Z41ehBCGouRhgojaRo,805
-ddapm_test_agent-1.37.0.dist-info/licenses/LICENSE.BSD3,sha256=J9S_Tq-hhvteDV2W8R0rqht5DZHkmvgdx3gnLZg4j6Q,1493
-ddapm_test_agent-1.37.0.dist-info/licenses/LICENSE.apache2,sha256=5V2RruBHZQIcPyceiv51DjjvdvhgsgS4pnXAOHDuZkQ,11342
-ddapm_test_agent-1.37.0.dist-info/METADATA,sha256=fjIqI75_eycNgtUDW4eNyF-Av0ybOJzXMImNk0-5q54,30465
-ddapm_test_agent-1.37.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ddapm_test_agent-1.37.0.dist-info/entry_points.txt,sha256=ulayVs6YJ-0Ej2kxbwn39wOHDVXbyQgFgsbRQmXydcs,250
-ddapm_test_agent-1.37.0.dist-info/top_level.txt,sha256=A9jiKOrrg6VjFAk-mtlSVYN4wr0VsZe58ehGR6IW47U,17
-ddapm_test_agent-1.37.0.dist-info/RECORD,,
+ddapm_test_agent-1.38.0.dist-info/licenses/LICENSE.BSD3,sha256=J9S_Tq-hhvteDV2W8R0rqht5DZHkmvgdx3gnLZg4j6Q,1493
+ddapm_test_agent-1.38.0.dist-info/licenses/LICENSE.apache2,sha256=5V2RruBHZQIcPyceiv51DjjvdvhgsgS4pnXAOHDuZkQ,11342
+ddapm_test_agent-1.38.0.dist-info/METADATA,sha256=HD007tJKa_D9iaGptUMvkjj8CN4zlH5_Xr78VYv0v5g,30466
+ddapm_test_agent-1.38.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ddapm_test_agent-1.38.0.dist-info/entry_points.txt,sha256=ulayVs6YJ-0Ej2kxbwn39wOHDVXbyQgFgsbRQmXydcs,250
+ddapm_test_agent-1.38.0.dist-info/top_level.txt,sha256=A9jiKOrrg6VjFAk-mtlSVYN4wr0VsZe58ehGR6IW47U,17
+ddapm_test_agent-1.38.0.dist-info/RECORD,,