feat(vcr): add ability to add custom providers in flag/environment variable (#250)

* add arg parse support for vcr provider map reading

* add vcr support for parsing custom providers map

* rel note

* add tests

* default fixture in conftest for vcr provider map

* rename test variables

* update readme with providers

* read custom ignore-list of headers for recording vcr cassettes

* use aiohttp_server instead

* add types pyyaml to riotfile for mypy

* safer get_custom_providers implementation

* add pyyaml to test deps

Author: sabrenner

README.md

@@ -148,7 +148,33 @@ The cassettes are matched based on the path, method, and body of the request. To
             -v $PWD/vcr-cassettes:/vcr-cassettes
             ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:latest
 
-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.
+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, DeepSeek, Anthropic, Google GenAI, and AWS Bedrock Runtime.
+
+#### Custom 3rd Party Providers
+
+The test agent can be configured to also register custom 3rd party providers. This is done by setting the `VCR_PROVIDER_MAP` environment variable or the `--vcr-provider-map` command-line option to a comma-separated list of provider names and their corresponding base URLs.
+
+```shell
+VCR_PROVIDER_MAP="provider1=http://provider1.com/,provider2=http://provider2.com/"
+```
+
+or
+
+```shell
+--vcr-provider-map="provider1=http://provider1.com/,provider2=http://provider2.com/"
+```
+
+The provider names are used to match the provider name in the request path, and the base URLs are used to proxy the request to the corresponding provider API endpoint.
+
+With this configuration set, you can make the following request to the test agent without error:
+
+```shell
+curl -X POST 'http://127.0.0.1:9126/vcr/provider1/some/path'
+```
+
+#### Ignoring Headers in Recorded Cassettes
+
+To ignore headers in recorded cassettes, you can use the `--vcr-ignore-headers` flag or `VCR_IGNORE_HEADERS` environment variable. The list should take the form of `header1,header2,header3`, and will be omitted from the recorded cassettes.
 
 #### 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.

ddapm_test_agent/agent.py

@@ -1580,6 +1580,8 @@ def make_app(
     snapshot_regex_placeholders: Dict[str, str],
     vcr_cassettes_directory: str,
     vcr_ci_mode: bool,
+    vcr_provider_map: str,
+    vcr_ignore_headers: str,
     enable_web_ui: bool = False,
 ) -> web.Application:
     agent = Agent()
@@ -1653,7 +1655,9 @@ def make_app(
             web.route(
                 "*",
                 "/vcr/{path:.*}",
-                lambda request: proxy_request(request, vcr_cassettes_directory, vcr_ci_mode),
+                lambda request: proxy_request(
+                    request, vcr_cassettes_directory, vcr_ci_mode, vcr_provider_map, vcr_ignore_headers
+                ),
             ),
         ]
     )
@@ -1961,6 +1965,18 @@ def main(args: Optional[List[str]] = None) -> None:
         default=os.environ.get("VCR_CI_MODE", False),
         help="Will change the test agent to record VCR cassettes in CI mode, throwing an error if a cassette is not found on /vcr/{provider}",
     )
+    parser.add_argument(
+        "--vcr-provider-map",
+        type=str,
+        default=os.environ.get("VCR_PROVIDER_MAP", ""),
+        help="Comma-separated list of provider=base_url tuples to map providers to paths. Used in addition to the default provider paths.",
+    )
+    parser.add_argument(
+        "--vcr-ignore-headers",
+        type=str,
+        default=os.environ.get("VCR_IGNORE_HEADERS", ""),
+        help="Comma-separated list of headers to ignore when recording VCR cassettes.",
+    )
     parser.add_argument(
         "--web-ui-port",
         type=int,
@@ -2017,6 +2033,8 @@ def main(args: Optional[List[str]] = None) -> None:
         snapshot_regex_placeholders=parsed_args.snapshot_regex_placeholders,
         vcr_cassettes_directory=parsed_args.vcr_cassettes_directory,
         vcr_ci_mode=parsed_args.vcr_ci_mode,
+        vcr_provider_map=parsed_args.vcr_provider_map,
+        vcr_ignore_headers=parsed_args.vcr_ignore_headers,
         enable_web_ui=parsed_args.web_ui_port > 0,
     )
 

ddapm_test_agent/vcr_proxy.py

@@ -1,3 +1,4 @@
+import asyncio
 import hashlib
 import json
 import logging
@@ -75,6 +76,16 @@ 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]:
+    return dict(
+        [
+            vcr_provider_map.strip().split("=", 1)
+            for vcr_provider_map in vcr_provider_map.split(",")
+            if vcr_provider_map.strip()
+        ]
+    )
+
+
 def normalize_multipart_body(body: bytes) -> str:
     if not body:
         return ""
@@ -114,14 +125,15 @@ def parse_authorization_header(auth_header: str) -> Dict[str, str]:
     return parsed
 
 
-def get_vcr(subdirectory: str, vcr_cassettes_directory: str) -> vcr.VCR:
+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,
+        filter_headers=CASSETTE_FILTER_HEADERS + extra_ignore_headers,
     )
 
 
@@ -146,7 +158,12 @@ def generate_cassette_name(path: str, method: str, body: bytes, vcr_cassette_pre
     )
 
 
-async def proxy_request(request: Request, vcr_cassettes_directory: str, vcr_ci_mode: bool) -> 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))
+
     path = request.match_info["path"]
     if request.query_string:
         path = path + "?" + request.query_string
@@ -156,7 +173,7 @@ async def proxy_request(request: Request, vcr_cassettes_directory: str, vcr_ci_m
         return Response(body="Invalid path format. Expected /{provider}/...", status=400)
 
     provider, remaining_path = parts
-    if provider not in PROVIDER_BASE_URLS:
+    if provider not in provider_base_urls:
         return Response(body=f"Unsupported provider: {provider}", status=400)
 
     body_bytes = await request.read()
@@ -173,7 +190,7 @@ async def proxy_request(request: Request, vcr_cassettes_directory: str, vcr_ci_m
             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] = {
@@ -200,8 +217,11 @@ async def proxy_request(request: Request, vcr_cassettes_directory: str, vcr_ci_m
         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)
+    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", "")

releasenotes/notes/vcr-custom-providers-map-a43b7e2d62d02015.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    vcr: Adds support for specifying a list of custom providers in the ``--vcr-provider-map`` flag or ``VCR_PROVIDER_MAP`` environment variable. The list should take the form of ``provider1=http://provider1.com/,provider2=http://provider2.com/``.
+  - |
+    vcr: Adds support for specifying a list of headers to ignore when recording VCR cassettes in the ``--vcr-ignore-headers`` flag or ``VCR_IGNORE_HEADERS`` environment variable. The list should take the form of ``header1,header2,header3``.

riotfile.py

@@ -63,6 +63,7 @@
                 "types-protobuf": latest,
                 "types-requests": latest,
                 "types-setuptools": latest,
+                "types-PyYAML": latest,
             },
         ),
         Venv(

test_deps.txt

@@ -1,3 +1,4 @@
 ddtrace==3.11.0
 pytest
 riot==0.20.1
+PyYAML==6.0.3

tests/conftest.py

@@ -106,15 +106,32 @@ def snapshot_regex_placeholders() -> Generator[Dict[str, str], None, None]:
 
 
 @pytest.fixture
-def snapshot_server_cassettes_directory() -> Generator[str, None, None]:
-    yield "/snapshot-server-cassettes"
+def vcr_cassettes_directory() -> Generator[str, None, None]:
+    import shutil
+    import tempfile
+
+    vcr_dir = tempfile.mkdtemp(prefix="vcr-cassettes-")
+    try:
+        yield vcr_dir
+    finally:
+        shutil.rmtree(vcr_dir, ignore_errors=True)
 
 
 @pytest.fixture
 def vcr_ci_mode() -> Generator[bool, None, None]:
     yield False
 
 
+@pytest.fixture
+def vcr_provider_map() -> Generator[str, None, None]:
+    yield ""
+
+
+@pytest.fixture
+def vcr_ignore_headers() -> Generator[str, None, None]:
+    yield ""
+
+
 @pytest.fixture
 async def agent_app(
     aiohttp_server,
@@ -130,25 +147,29 @@ async def agent_app(
     disable_error_responses,
     snapshot_removed_attrs,
     snapshot_regex_placeholders,
-    snapshot_server_cassettes_directory,
+    vcr_cassettes_directory,
     vcr_ci_mode,
+    vcr_provider_map,
+    vcr_ignore_headers,
 ):
     app = await aiohttp_server(
         make_app(
-            agent_enabled_checks,
-            log_span_fmt,
-            str(snapshot_dir),
-            snapshot_ci_mode,
-            snapshot_ignored_attrs,
-            agent_url,
-            trace_request_delay,
-            suppress_trace_parse_errors,
-            pool_trace_check_failures,
-            disable_error_responses,
-            snapshot_removed_attrs,
-            snapshot_regex_placeholders,
-            snapshot_server_cassettes_directory,
-            vcr_ci_mode,
+            enabled_checks=agent_enabled_checks,
+            log_span_fmt=log_span_fmt,
+            snapshot_dir=str(snapshot_dir),
+            snapshot_ci_mode=snapshot_ci_mode,
+            snapshot_ignored_attrs=snapshot_ignored_attrs,
+            agent_url=agent_url,
+            trace_request_delay=trace_request_delay,
+            suppress_trace_parse_errors=suppress_trace_parse_errors,
+            pool_trace_check_failures=pool_trace_check_failures,
+            disable_error_responses=disable_error_responses,
+            snapshot_removed_attrs=snapshot_removed_attrs,
+            snapshot_regex_placeholders=snapshot_regex_placeholders,
+            vcr_cassettes_directory=vcr_cassettes_directory,
+            vcr_ci_mode=vcr_ci_mode,
+            vcr_provider_map=vcr_provider_map,
+            vcr_ignore_headers=vcr_ignore_headers,
         )
     )
     yield app