Add snapshot commands (#96)

Add commands to start and perform snapshot test cases to ease
manual/CI usage. The alternative would be to curl and try to parse the
results which isn't trivial to do.


Co-authored-by: Brett Langdon <me@brett.is>

Author: Kyle-Verhoog

ddapm_test_agent/cmd.py

@@ -0,0 +1,78 @@
+import argparse
+import os
+import sys
+
+import requests
+import yarl
+
+
+def _add_token_arg(parser: argparse.ArgumentParser) -> None:
+    """Add the test session token argument to the parser."""
+    parser.add_argument(
+        "--test-session-token",
+        type=str,
+        default=os.environ.get("TEST_SESSION_TOKEN"),
+        help="Test session token to query with.",
+    )
+
+
+def _add_agent_url_arg(parser: argparse.ArgumentParser) -> None:
+    """Add the agent url argument to the given parser."""
+    parser.add_argument(
+        "--agent-url",
+        type=str,
+        default=os.environ.get(
+            "DD_TRACE_AGENT_URL",
+            os.environ.get("DD_AGENT_URL", "http://localhost:8126"),
+        ),
+        help=("Test agent URL. Default is http://localhost:8126"),
+    )
+
+
+def main_session_start() -> None:
+    """Entrypoint for the start-session command"""
+    parser = argparse.ArgumentParser(
+        description=(
+            "Start a test agent session with a given token. "
+            "All data submitted with this token will be associated with this session."
+        ),
+        prog="ddapm-test-agent-start-session",
+    )
+    _add_agent_url_arg(parser)
+    _add_token_arg(parser)
+    parsed_args = parser.parse_args(sys.argv[1:])
+    url = yarl.URL(parsed_args.agent_url).with_path("/test/session/start")
+    resp = requests.get(
+        str(url), params={"test_session_token": parsed_args.test_session_token}
+    )
+    if resp.status_code != 200:
+        print(resp.text)
+        sys.exit(1)
+    print(resp.text)
+    sys.exit(0)
+
+
+def main_snapshot() -> None:
+    """Entrypoint for the snapshot command"""
+    parser = argparse.ArgumentParser(
+        description=("Perform a snapshot test for the data received in the session."),
+        prog="ddapm-test-agent-snapshot",
+    )
+    _add_agent_url_arg(parser)
+    _add_token_arg(parser)
+    parsed_args = parser.parse_args(sys.argv[1:])
+    if not parsed_args.test_session_token:
+        print(
+            "Error: a test token is required! Please specify one with --test-session-token"
+            " command line argument or the TEST_SESSION_TOKEN environment variable."
+        )
+        sys.exit(1)
+    url = yarl.URL(parsed_args.agent_url).with_path("/test/session/snapshot")
+    resp = requests.get(
+        str(url), params={"test_session_token": parsed_args.test_session_token}
+    )
+    if resp.status_code != 200:
+        print(resp.text)
+        sys.exit(1)
+    print(resp.text)
+    sys.exit(0)

releasenotes/notes/snapshot-cmds-de4a561911afee11.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Add ``ddapm-test-agent-session-start`` and ``ddapm-test-agent-snapshot`` commands
+    to provide a more convenient way to perform snapshots manually or in CI.

riotfile.py

@@ -57,8 +57,9 @@
             pkgs={
                 "mypy": latest,
                 "pytest": latest,
-                "types-setuptools": latest,
                 "types-protobuf": latest,
+                "types-requests": latest,
+                "types-setuptools": latest,
             },
         ),
         Venv(

setup.py

@@ -30,7 +30,9 @@
         "aiohttp",
         "ddsketch",
         "msgpack",
+        "requests",
         "typing_extensions",
+        "yarl",
     ],
     tests_require=testing_deps,
     setup_requires=["setuptools_scm"],
@@ -39,6 +41,8 @@
         "console_scripts": [
             "ddapm-test-agent=ddapm_test_agent.agent:main",
             "ddapm-test-agent-fmt=ddapm_test_agent.fmt:main",
+            "ddapm-test-agent-session-start=ddapm_test_agent.cmd:main_session_start",
+            "ddapm-test-agent-snapshot=ddapm_test_agent.cmd:main_snapshot",
         ]
     },
     extras_require={

tests/test_snapshot_integration.py

@@ -321,3 +321,41 @@ async def test_tracestats(
         assert resp.status == 400
     else:
         assert resp.status == 200
+
+
+async def test_cmd(testagent: aiohttp.ClientSession, tracer: Tracer) -> None:
+    """Test the commands provided with the library.
+
+    Note that this test reuses the trace/snapshot from test_single_trace above.
+    """
+    env = os.environ.copy()
+    p = subprocess.run(
+        ["ddapm-test-agent-session-start", "--test-session-token=test_single_trace"],
+        env=env,
+    )
+    assert p.returncode == 0
+    with tracer.trace(
+        "root", service="custom_service", resource="/url/endpoint", span_type="web"
+    ):
+        pass
+    tracer.flush()
+    p = subprocess.run(
+        ["ddapm-test-agent-snapshot", "--test-session-token=test_single_trace"], env=env
+    )
+    assert p.returncode == 0
+
+    # Ensure failing snapshots work.
+    p = subprocess.run(
+        ["ddapm-test-agent-session-start", "--test-session-token=test_single_trace"],
+        env=env,
+    )
+    assert p.returncode == 0
+    with tracer.trace(
+        "root1234", service="custom_service", resource="/url/endpoint", span_type="web"
+    ):
+        pass
+    tracer.shutdown()
+    p = subprocess.run(
+        ["ddapm-test-agent-snapshot", "--test-session-token=test_single_trace"], env=env
+    )
+    assert p.returncode == 1