Snapshot record then reply testing

Author: mateusz834

.github/workflows/test.yml

@@ -19,11 +19,11 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
           deps-group: test
-      - name: Download Splunk MCP Server App
-        run: uv run ./scripts/download_splunk_mcp_server_app.py
-        env:
-          SPLUNKBASE_USERNAME: ${{ secrets.SPLUNKBASE_USERNAME }}
-          SPLUNKBASE_PASSWORD: ${{ secrets.SPLUNKBASE_PASSWORD }}
+      #- name: Download Splunk MCP Server App
+      #  run: uv run ./scripts/download_splunk_mcp_server_app.py
+      #  env:
+      #    SPLUNKBASE_USERNAME: ${{ secrets.SPLUNKBASE_USERNAME }}
+      #    SPLUNKBASE_PASSWORD: ${{ secrets.SPLUNKBASE_PASSWORD }}
       - name: Launch Splunk Docker instance
         run: SPLUNK_VERSION=${{ matrix.splunk-version }} docker compose up -d
       - name: Set up .env

pyproject.toml

@@ -45,6 +45,7 @@ test = [
   "pytest-cov>=7.1.0",
   "pytest-asyncio>=1.3.0",
   "python-dotenv>=1.2.2",
+  "pytest-recording>=0.13.4",
 ]
 release = ["build>=1.4.3", "jinja2>=3.1.6", "sphinx>=9.1.0", "twine>=6.2.0"]
 lint = ["basedpyright>=1.39.0", "ruff>=0.15.10"]

tests/ai_testlib.py

@@ -1,4 +1,15 @@
-from typing import override
+import functools
+import inspect
+import json
+import os
+from collections.abc import Callable, Coroutine
+from typing import Any, override
+from urllib import parse
+
+import vcr
+from vcr.config import RecordMode
+from vcr.request import Request
+
 from splunklib.ai.model import PredefinedModel
 from tests.ai_test_model import InternalAIModel, TestLLMSettings, create_model
 from tests.testlib import SDKTestCase
@@ -42,3 +53,121 @@ async def model(self) -> PredefinedModel:
         model = await create_model(self.test_llm_settings)
         self._model = model
         return model
+
+
+def ai_snapshot_test() -> Callable[
+    [Callable[..., Coroutine[Any, Any, None]]], Callable[..., Coroutine[Any, Any, None]]
+]:
+    def decorator(
+        fn: Callable[..., Coroutine[Any, Any, None]],
+    ) -> Callable[..., Coroutine[Any, Any, None]]:
+        source_file = inspect.getfile(fn)
+        test_dir = os.path.dirname(source_file)
+        test_file = os.path.splitext(os.path.basename(source_file))[0]
+
+        snapshot_path = os.path.join(
+            test_dir, "snapshots", test_file, f"{fn.__qualname__}.json"
+        )
+
+        # TODO: DROP THIS, and use the dir in VCR.
+        os.makedirs(os.path.dirname(snapshot_path), exist_ok=True)
+
+        @functools.wraps(fn)
+        async def wrapper(self: AITestCase, *args: Any, **kwargs: Any) -> None:
+            settings = self.test_llm_settings
+            assert settings.internal_ai is not None
+
+            internal_ai_hostname = parse.urlparse(
+                settings.internal_ai.base_url
+            ).hostname
+
+            REDACTED_APP_KEY = "[[[--APPKEY-REDACTED-]]]"
+
+            class _JSONFriendlySerializer:
+                def deserialize(self, serialized: str) -> Any:
+                    assert settings.internal_ai is not None
+                    serialized = serialized.replace(
+                        REDACTED_APP_KEY, settings.internal_ai.app_key
+                    )
+
+                    data = json.loads(serialized)
+                    for interaction in data.get("interactions", []):
+                        interaction["request"]["uri"] = interaction["request"][
+                            "uri"
+                        ].replace("internal-ai-host", internal_ai_hostname, 1)
+
+                        interaction["request"]["body"] = json.dumps(
+                            interaction["request"]["body"]
+                        )
+                        body = interaction["response"]["body"]
+                        interaction["response"]["body"] = {}
+                        interaction["response"]["body"]["string"] = json.dumps(body)
+
+                    return data
+
+                def serialize(self, dict: Any) -> str:
+                    for interaction in dict.get("interactions", []):
+                        interaction["request"]["uri"] = interaction["request"][
+                            "uri"
+                        ].replace(internal_ai_hostname, "internal-ai-host", 1)
+
+                        body = interaction["request"]["body"]
+                        interaction["request"]["body"] = json.loads(body)
+
+                        resp_body = interaction["response"]["body"]["string"]
+                        interaction["response"]["body"] = json.loads(resp_body)
+
+                    # TODO: redact app_key with [KEY-REDACTED-KAJFKDJFKSDJKFUSFIUIUR]
+                    # TODO: assert that nothing from settings is in the dump.
+                    out = json.dumps(dict, indent=4) + "\n"
+                    assert settings.internal_ai is not None
+                    out = out.replace(settings.internal_ai.app_key, REDACTED_APP_KEY)
+
+                    # Assert
+
+                    return out
+
+            def _before_record_request(request: Request) -> Request | None:
+                url = parse.urlparse(request.uri)
+                if url.hostname == internal_ai_hostname:
+                    request.headers = {}
+                    return request
+                return None
+
+            def _before_record_response(response: Any) -> Any:
+                response["headers"] = {}
+                return response
+
+            def _json_body_matcher(r1: Any, r2: Any) -> None:
+                b1 = json.loads(r1.body)
+                b2 = json.loads(r2.body)
+                if b1 != b2:
+                    raise AssertionError(f"Body mismatch:\n{b1}\n!=\n{b2}")
+
+            my_vcr = vcr.VCR(
+                cassette_library_dir=".",
+                serializer="json-friendly",
+                record_mode=RecordMode.ONCE,
+                match_on=[
+                    "method",
+                    "scheme",
+                    "host",
+                    "port",
+                    "path",
+                    "query",
+                    "jsonbody",
+                ],
+                before_record_request=_before_record_request,
+                before_record_response=_before_record_response,
+                # record_on_exception=False,
+                drop_unused_requests=True,
+            )
+            my_vcr.register_serializer("json-friendly", _JSONFriendlySerializer())
+            my_vcr.register_matcher("jsonbody", _json_body_matcher)
+
+            with my_vcr.use_cassette(snapshot_path):
+                await fn(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator