ddapm-test-agent 1.32.0__py3-none-any.whl → 1.33.1__py3-none-any.whl

ddapm_test_agent/__init__.py

@@ -1,4 +1,4 @@
 def _get_version():
-    import pkg_resources  # type: ignore[import-untyped]
+    import pkg_resources
 
     return pkg_resources.get_distribution(__name__).version

ddapm_test_agent/agent.py

@@ -133,7 +133,7 @@ def _session_token(request: Request) -> Optional[str]:
 
 async def _vcr_proxy_cassette_prefix(request: Request) -> Optional[str]:
     try:
-        request_body: dict[str, str] = await request.json()
+        request_body: Dict[str, str] = await request.json()
         requested_test_name = request_body.get("test_name")
         return requested_test_name
     except (json.JSONDecodeError, UnicodeDecodeError):
@@ -165,7 +165,7 @@ async def handle_exception_middleware(request: Request, handler: _Handler) -> we
 
 async def _forward_request(
     request_data: bytes, headers: Mapping[str, str], full_agent_url: str
-) -> tuple[ClientResponse, str]:
+) -> Tuple[ClientResponse, str]:
     async with ClientSession() as session:
         async with session.post(
             full_agent_url,
@@ -929,10 +929,9 @@ class Agent:
         # Get the span attributes that are to be removed for this snapshot.
         default_attribute_regex_replaces: Dict[str, str] = request.app["snapshot_regex_placeholders"]
         regex_overrides = _parse_map(request.url.query.get("regex_placeholders", ""))
-        attribute_regex_replaces = dict(
-            (f"{{{key}}}", re.compile(regex))
-            for (key, regex) in (default_attribute_regex_replaces | regex_overrides).items()
-        )
+        regex_replaces = default_attribute_regex_replaces.copy()
+        regex_replaces.update(regex_overrides)
+        attribute_regex_replaces = dict((f"{{{key}}}", re.compile(regex)) for (key, regex) in regex_replaces.items())
         log.info("using regex placeholders %r", attribute_regex_replaces)
 
         if "span_id" in span_removes:

ddapm_test_agent/trace_snapshot.py

@@ -3,13 +3,13 @@ import json
 import logging
 import operator
 import pprint
-from re import Pattern
 import textwrap
 from typing import Any
 from typing import Dict
 from typing import List
 from typing import Optional
 from typing import OrderedDict as OrderedDictType
+from typing import Pattern
 from typing import Set
 from typing import Tuple
 from typing import cast
@@ -188,9 +188,13 @@ def _match_traces(t1s: List[Trace], t2s: List[Trace]) -> List[Tuple[Trace, Trace
     return matches
 
 
-def _diff_spans(
-    s1: Span, s2: Span, ignored: Set[str]
-) -> Tuple[List[str], List[str], List[str], List[Tuple[List[str], List[str]]], List[Tuple[List[str], List[str]]]]:
+def _diff_spans(s1: Span, s2: Span, ignored: Set[str]) -> Tuple[
+    List[str],
+    List[str],
+    List[str],
+    List[Tuple[List[str], List[str]]],
+    List[Tuple[List[str], List[str]]],
+]:
     """Return differing attributes between two spans and their meta/metrics maps.
 
     It is assumed that the spans have passed through preliminary validation
@@ -255,7 +259,11 @@ def _diff_spans(
             )
             link_diff = []
             for d1, d2, ign in [
-                (l1_no_tags, l2_no_tags, set(i[11:] for i in ignored if i.startswith("span_links."))),
+                (
+                    l1_no_tags,
+                    l2_no_tags,
+                    set(i[11:] for i in ignored if i.startswith("span_links.")),
+                ),
                 (
                     l1.get("attributes") or {},
                     l2.get("attributes") or {},
@@ -288,7 +296,11 @@ def _diff_spans(
             )
             event_diff = []
             for d1, d2, ign in [
-                (l1_no_tags, l2_no_tags, set(i[12:] for i in ignored if i.startswith("span_events."))),
+                (
+                    l1_no_tags,
+                    l2_no_tags,
+                    set(i[12:] for i in ignored if i.startswith("span_events.")),
+                ),
                 (
                     e1.get("attributes") or {},
                     e2.get("attributes") or {},
@@ -307,7 +319,13 @@ def _diff_spans(
     results.append(event_diffs)  # type: ignore
 
     return cast(
-        Tuple[List[str], List[str], List[str], List[Tuple[List[str], List[str]]], List[Tuple[List[str], List[str]]]],
+        Tuple[
+            List[str],
+            List[str],
+            List[str],
+            List[Tuple[List[str], List[str]]],
+            List[Tuple[List[str], List[str]]],
+        ],
         tuple(results),
     )
 
@@ -334,9 +352,13 @@ def _compare_traces(expected: Trace, received: Trace, ignored: Set[str]) -> None
         ) as frame:
             frame.add_item(f"Expected span:\n{pprint.pformat(s_exp)}")
             frame.add_item(f"Received span:\n{pprint.pformat(s_rec)}")
-            top_level_diffs, meta_diffs, metrics_diffs, span_link_diffs, span_event_diffs = _diff_spans(
-                s_exp, s_rec, ignored
-            )
+            (
+                top_level_diffs,
+                meta_diffs,
+                metrics_diffs,
+                span_link_diffs,
+                span_event_diffs,
+            ) = _diff_spans(s_exp, s_rec, ignored)
 
             for diffs, diff_type, d_exp, d_rec in [
                 (top_level_diffs, "span", s_exp, s_rec),
@@ -363,7 +385,12 @@ def _compare_traces(expected: Trace, received: Trace, ignored: Set[str]) -> None
 
             for i, (link_level_diffs, attribute_diffs) in enumerate(span_link_diffs):
                 for diffs, diff_type, d_exp, d_rec in [
-                    (link_level_diffs, f"{i}", s_exp["span_links"][i], s_rec["span_links"][i]),
+                    (
+                        link_level_diffs,
+                        f"{i}",
+                        s_exp["span_links"][i],
+                        s_rec["span_links"][i],
+                    ),
                     (
                         attribute_diffs,
                         f"{i} attributes",
@@ -387,7 +414,12 @@ def _compare_traces(expected: Trace, received: Trace, ignored: Set[str]) -> None
 
             for i, (event_level_diffs, attribute_diffs) in enumerate(span_event_diffs):
                 for diffs, diff_type, d_exp, d_rec in [
-                    (event_level_diffs, f"{i}", s_exp["span_events"][i], s_rec["span_events"][i]),
+                    (
+                        event_level_diffs,
+                        f"{i}",
+                        s_exp["span_events"][i],
+                        s_rec["span_events"][i],
+                    ),
                     (
                         attribute_diffs,
                         f"{i} attributes",
@@ -545,4 +577,7 @@ def generate_snapshot(
     removed: Optional[List[str]] = None,
     attribute_regex_replaces: Optional[Dict[str, Pattern[str]]] = None,
 ) -> str:
-    return _snapshot_json(_normalize_traces(received_traces, attribute_regex_replaces or {}), removed or [])
+    return _snapshot_json(
+        _normalize_traces(received_traces, attribute_regex_replaces or {}),
+        removed or [],
+    )

ddapm_test_agent/vcr_proxy.py

@@ -1,21 +1,38 @@
 import hashlib
 import json
+import logging
 import os
 import re
+from typing import Any
+from typing import Dict
 from typing import Optional
 from urllib.parse import urljoin
 
 from aiohttp.web import Request
 from aiohttp.web import Response
 import requests
+from requests_aws4auth import AWS4Auth
 import vcr
 
 
+logger = logging.getLogger(__name__)
+
+
+#  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",
+}
+
+
 PROVIDER_BASE_URLS = {
     "openai": "https://api.openai.com/v1",
     "azure-openai": "https://dd.openai.azure.com/",
@@ -23,8 +40,25 @@ PROVIDER_BASE_URLS = {
     "anthropic": "https://api.anthropic.com/",
     "datadog": "https://api.datadoghq.com/",
     "genai": "https://generativelanguage.googleapis.com/",
+    "bedrock-runtime": f"https://bedrock-runtime.{AWS_REGION}.amazonaws.com",
 }
 
+CASSETTE_FILTER_HEADERS = [
+    "authorization",
+    "OpenAI-Organization",
+    "api-key",
+    "x-api-key",
+    "dd-api-key",
+    "dd-application-key",
+    "x-goog-api-key",
+    "x-amz-security-token",
+    "x-amz-content-sha256",
+    "x-amz-date",
+    "x-amz-user-agent",
+    "amz-sdk-invocation-id",
+    "amz-sdk-request",
+]
+
 NORMALIZERS = [
     (
         r"--form-data-boundary-[^\r\n]+",
@@ -65,6 +99,21 @@ def normalize_multipart_body(body: bytes) -> str:
             return f"[binary_data_{hex_digest}]"
 
 
+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 {}
+
+    auth_parts = auth_header[len("AWS4-HMAC-SHA256 ") :].split(",")
+    parsed = {}
+
+    for part in auth_parts:
+        key, value = part.split("=", 1)
+        parsed[key.strip()] = value.strip()
+
+    return parsed
+
+
 def get_vcr(subdirectory: str, vcr_cassettes_directory: str) -> vcr.VCR:
     cassette_dir = os.path.join(vcr_cassettes_directory, subdirectory)
 
@@ -72,15 +121,7 @@ def get_vcr(subdirectory: str, vcr_cassettes_directory: str) -> vcr.VCR:
         cassette_library_dir=cassette_dir,
         record_mode="once",
         match_on=["path", "method"],
-        filter_headers=[
-            "authorization",
-            "OpenAI-Organization",
-            "api-key",
-            "x-api-key",
-            "dd-api-key",
-            "dd-application-key",
-            "x-goog-api-key",
-        ],
+        filter_headers=CASSETTE_FILTER_HEADERS,
     )
 
 
@@ -125,31 +166,50 @@ async def proxy_request(request: Request, vcr_cassettes_directory: str) -> Respo
     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)
-    with get_vcr(provider, vcr_cassettes_directory).use_cassette(f"{cassette_name}.yaml"):
-        oai_response = requests.request(
-            method=request.method,
-            url=target_url,
-            headers=headers,
-            data=body_bytes,
-            cookies=dict(request.cookies),
-            allow_redirects=False,
-            stream=True,
-        )
+    cassette_file_name = f"{cassette_name}.yaml"
+
+    request_kwargs: Dict[str, Any] = {
+        "method": request.method,
+        "url": target_url,
+        "headers": headers,
+        "data": body_bytes,
+        "cookies": dict(request.cookies),
+        "allow_redirects": False,
+        "stream": True,
+    }
+
+    if provider in AWS_SERVICES and not os.path.exists(
+        os.path.join(vcr_cassettes_directory, provider, cassette_file_name)
+    ):
+        if not AWS_SECRET_ACCESS_KEY:
+            return Response(
+                body="AWS_SECRET_ACCESS_KEY environment variable not set for aws signature recalculation",
+                status=400,
+            )
+
+        auth_header = request.headers.get("Authorization", "")
+        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
+
+    with get_vcr(provider, vcr_cassettes_directory).use_cassette(cassette_file_name):
+        provider_response = requests.request(**request_kwargs)
 
     # Extract content type without charset
-    content_type = oai_response.headers.get("content-type", "")
+    content_type = provider_response.headers.get("content-type", "")
     if ";" in content_type:
         content_type = content_type.split(";")[0].strip()
 
     response = Response(
-        body=oai_response.content,
-        status=oai_response.status_code,
+        body=provider_response.content,
+        status=provider_response.status_code,
         content_type=content_type,
     )
 
-    for key, value in oai_response.headers.items():
+    for key, value in provider_response.headers.items():
         if key.lower() not in (
             "content-length",
             "transfer-encoding",

{ddapm_test_agent-1.32.0.dist-info → ddapm_test_agent-1.33.1.dist-info}/METADATA

@@ -1,15 +1,12 @@
 Metadata-Version: 2.4
 Name: ddapm-test-agent
-Version: 1.32.0
+Version: 1.33.1
 Summary: Test agent for Datadog APM client libraries
 Home-page: https://github.com/Datadog/dd-apm-test-agent
 Author: Kyle Verhoog
 Author-email: kyle@verhoog.ca
 License: BSD 3
 Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
@@ -24,6 +21,7 @@ Requires-Dist: requests
 Requires-Dist: typing_extensions
 Requires-Dist: yarl
 Requires-Dist: vcrpy
+Requires-Dist: requests-aws4auth
 Requires-Dist: opentelemetry-proto<1.37.0,>1.33.0
 Requires-Dist: protobuf>=3.19.0
 Requires-Dist: grpcio<2.0,>=1.66.2
@@ -63,11 +61,6 @@ See the [Development](#development) section for how to get the test agent runnin
 
 ## Installation
 
-The test agent can be installed using [nix](https://docs.determinate.systems/getting-started/):
-
-    nix profile install github:datadog/dd-apm-test-agent#ddapm-test-agent
-    # nix profile upgrade ddapm-test-agent # to upgrade
-
 The test agent can be installed from PyPI:
 
     pip install ddapm-test-agent
@@ -192,6 +185,13 @@ The cassettes are matched based on the path, method, and body of the request. To
 
 Optionally specifying whatever mounted path is used for the cassettes directory. The test agent comes with a default set of cassettes for OpenAI, Azure OpenAI, and DeepSeek.
 
+#### AWS Services
+AWS service proxying, specifically recording cassettes for the first time, requires a `AWS_SECRET_ACCESS_KEY` environment variable to be set for the container running the test agent. This is used to recalculate the AWS signature for the request, as the one generated client-side likely used `{test-agent-host}:{test-agent-port}/vcr/{aws-service}` as the host, and the signature will mismatch that on the actual AWS service.
+
+Additionally, the `AWS_REGION` environment variable can be set, defaulting to `us-east-1`.
+
+To add a new AWS service to proxy, add an entry in the `PROVIDER_BASE_URLS` for its provider url, and an entry in the `AWS_SERVICES` dictionary for the service name, since they are not always a one-to-one mapping with the implied provider url (e.g, `https://bedrock-runtime.{AWS_REGION}.amazonaws.com` is the provider url, but the service name is `bedrock`, as `bedrock` also has multiple sub services, like `converse`).
+
 #### Usage in clients
 
 To use this feature in your client, you can use the `/vcr/{provider}` endpoint to proxy requests to the provider API.

{ddapm_test_agent-1.32.0.dist-info → ddapm_test_agent-1.33.1.dist-info}/RECORD

@@ -1,5 +1,5 @@
-ddapm_test_agent/__init__.py,sha256=hJBQduemued6MAWH6Uk2gqJx2vb0_dd0UO52pJ83nLM,138
-ddapm_test_agent/agent.py,sha256=S5VzW5yvBQKqsGj50Xqb_u8gXvRnZguwWTZXn6BcX-M,73960
+ddapm_test_agent/__init__.py,sha256=IEYMDM-xI0IoHYSYw4Eva5263puB_crrrbLstOCScRw,106
+ddapm_test_agent/agent.py,sha256=296wweDNsF7cw5QuGdSzbKa4vWrVpea3rvZU4Laddkc,74000
 ddapm_test_agent/apmtelemetry.py,sha256=w_9-yUDh1dgox-FfLqeOHU2C14GcjOjen-_SVagiZrc,861
 ddapm_test_agent/checks.py,sha256=pBa4YKZQVA8qaTVJ_XgMA6TmlUZNh99YOrCFJA7fwo0,6865
 ddapm_test_agent/client.py,sha256=ViEmiRX9Y3SQ-KBhSc-FdzBmIVIe8Ij9jj-Q6VGyzLY,7359
@@ -12,15 +12,15 @@ ddapm_test_agent/metrics.py,sha256=EZo7lSec2oAiH7tUqavKZ2MJM7TwbuFGE3AT3cXwmSM,3
 ddapm_test_agent/remoteconfig.py,sha256=_QjYUKc3JF31DxdvISDXgslm5WVnYWAw0hyckWuLc1c,3606
 ddapm_test_agent/trace.py,sha256=t0OR8w3NcZK-EOOoadgPITiZqS5tAJGtxqLVGLEw7Kg,45816
 ddapm_test_agent/trace_checks.py,sha256=bRg2eLKoHROXIFJRbujMUn0T3x1X8pZso-j8wXNomec,9972
-ddapm_test_agent/trace_snapshot.py,sha256=g2MhKi8UE-Wsf6PtuzPoXymcW-cYRUnvj63SP9FETJs,22354
+ddapm_test_agent/trace_snapshot.py,sha256=ayOUcCFo6xyotFRm0fSNdeA91_T447W5ShlkRvi0nZE,22932
 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=g6ix7laiS8Hqq9p14nkTMARhj5KMZmyRZjZpfFEMxOM,4973
-ddapm_test_agent-1.32.0.dist-info/licenses/LICENSE.BSD3,sha256=J9S_Tq-hhvteDV2W8R0rqht5DZHkmvgdx3gnLZg4j6Q,1493
-ddapm_test_agent-1.32.0.dist-info/licenses/LICENSE.apache2,sha256=5V2RruBHZQIcPyceiv51DjjvdvhgsgS4pnXAOHDuZkQ,11342
-ddapm_test_agent-1.32.0.dist-info/METADATA,sha256=mf922XgdhlbHNJk3P5jCiw5HBkmzfwpex-FDvqjHXa4,27493
-ddapm_test_agent-1.32.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ddapm_test_agent-1.32.0.dist-info/entry_points.txt,sha256=ulayVs6YJ-0Ej2kxbwn39wOHDVXbyQgFgsbRQmXydcs,250
-ddapm_test_agent-1.32.0.dist-info/top_level.txt,sha256=A9jiKOrrg6VjFAk-mtlSVYN4wr0VsZe58ehGR6IW47U,17
-ddapm_test_agent-1.32.0.dist-info/RECORD,,
+ddapm_test_agent/vcr_proxy.py,sha256=jTm9Q0LeWOUJrwlDrG632WrQB6ScWvjb3dLR7oSFa6o,6808
+ddapm_test_agent-1.33.1.dist-info/licenses/LICENSE.BSD3,sha256=J9S_Tq-hhvteDV2W8R0rqht5DZHkmvgdx3gnLZg4j6Q,1493
+ddapm_test_agent-1.33.1.dist-info/licenses/LICENSE.apache2,sha256=5V2RruBHZQIcPyceiv51DjjvdvhgsgS4pnXAOHDuZkQ,11342
+ddapm_test_agent-1.33.1.dist-info/METADATA,sha256=fPwXRZgry4efp28Ft0jEWhzxwW8rZx134I01GBFHt0c,28104
+ddapm_test_agent-1.33.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ddapm_test_agent-1.33.1.dist-info/entry_points.txt,sha256=ulayVs6YJ-0Ej2kxbwn39wOHDVXbyQgFgsbRQmXydcs,250
+ddapm_test_agent-1.33.1.dist-info/top_level.txt,sha256=A9jiKOrrg6VjFAk-mtlSVYN4wr0VsZe58ehGR6IW47U,17
+ddapm_test_agent-1.33.1.dist-info/RECORD,,