Add delay for telemetry requests (#198)

Extend --trace-request-delay to also affect trace requests.

Also, add a test/settings endpoint to change the value of settings on the fly.

README.md

@@ -439,6 +439,22 @@ Returns the tracer-flares in the following json format:
 
 If there was an error parsing the tracer-flare form, that will be recorded under `error`.
 
+### /test/settings (POST)
+
+Allows to change some settings on the fly.
+This endpoint takes a POST request with a json content listing the keys and values to apply.
+
+```json
+{ 'key': value }
+```
+
+Supported keys:
+ - `trace_request_delay`: sets a delay to apply to trace and telemetry requests
+
+ ```
+curl -X POST 'http://0.0.0.0:8126/test/settings' -d '{ "trace_request_delay": 5 }'
+```
+
 ## Development
 
 ### Prerequisites

ddapm_test_agent/agent.py

@@ -38,7 +38,7 @@
 from . import trace_snapshot
 from . import tracestats_snapshot
 from .apmtelemetry import TelemetryEvent
-from .apmtelemetry import v2_decode as v2_apmtelemetry_decode
+from .apmtelemetry import v2_decode_request as v2_apmtelemetry_decode_request
 from .checks import CheckTrace
 from .checks import Checks
 from .checks import start_trace
@@ -332,7 +332,7 @@ async def apmtelemetry(self) -> List[TelemetryEvent]:
         _events: List[TelemetryEvent] = []
         for req in reversed(self._requests):
             if req.match_info.handler == self.handle_v2_apmtelemetry:
-                _events.append(v2_apmtelemetry_decode(await req.read()))
+                _events.append(await v2_apmtelemetry_decode_request(req, await req.read()))
         return _events
 
     async def _trace_by_trace_id(self, trace_id: int) -> Trace:
@@ -435,7 +435,7 @@ async def _apmtelemetry_by_session(self, token: Optional[str]) -> List[Telemetry
         events: List[TelemetryEvent] = []
         for req in self._requests_by_session(token):
             if req.match_info.handler == self.handle_v2_apmtelemetry:
-                events.append(v2_apmtelemetry_decode(await req.read()))
+                events.append(await v2_apmtelemetry_decode_request(req, await req.read()))
 
         # TODO: Sort the events?
         return events
@@ -590,7 +590,7 @@ async def handle_v07_remoteconfig_put(self, request: Request) -> web.Response:
         return web.HTTPAccepted()
 
     async def handle_v2_apmtelemetry(self, request: Request) -> web.Response:
-        v2_apmtelemetry_decode(self._request_data(request))
+        await v2_apmtelemetry_decode_request(request, self._request_data(request))
         # TODO: Validation
         # TODO: Snapshots
         return web.HTTPOk()
@@ -661,6 +661,22 @@ async def handle_get_tested_integrations(self, request: Request) -> web.Response
             aggregated_text = ",".join(text_headers) + "\n" + aggregated_text
         return web.Response(body=aggregated_text, content_type="text/plain", headers=req_headers)
 
+    async def handle_settings(self, request: Request) -> web.Response:
+        """Allow to change test agent settings on the fly"""
+        raw_data = await request.read()
+        data = json.loads(raw_data)
+
+        # First pass to validate the data
+        for key in data:
+            if key not in request.app:
+                return web.HTTPUnprocessableEntity(text=f"Unknown key: '{key}'")
+
+        # Second pass to apply the config
+        for key in data:
+            request.app[key] = data[key]
+
+        return web.HTTPAccepted()
+
     async def handle_info(self, request: Request) -> web.Response:
         return web.json_response(
             {
@@ -1142,6 +1158,7 @@ def make_app(
             web.get("/test/trace_check/clear", agent.clear_trace_check_failures),
             web.get("/test/trace_check/summary", agent.get_trace_check_summary),
             web.get("/test/integrations/tested_versions", agent.handle_get_tested_integrations),
+            web.post("/test/settings", agent.handle_settings),
         ]
     )
     checks = Checks(
@@ -1254,7 +1271,7 @@ def main(args: Optional[List[str]] = None) -> None:
         "--trace-request-delay",
         type=float,
         default=os.environ.get("DD_TEST_STALL_REQUEST_SECONDS", 0.0),
-        help=("Will stall trace requests for specified amount of time"),
+        help=("Will stall trace and telemetry requests for specified amount of time"),
     )
     parser.add_argument(
         "--suppress-trace-parse-errors",

ddapm_test_agent/apmtelemetry.py

@@ -1,12 +1,30 @@
+import asyncio
 import json
+import logging
 from typing import Any
 from typing import Dict
 from typing import cast
 
+from aiohttp.web import Request
 
+
+log = logging.getLogger(__name__)
 TelemetryEvent = Dict[str, Any]
 
 
+async def v2_decode_request(request: Request, data: bytes) -> TelemetryEvent:
+    headers = request.headers
+
+    if "X-Datadog-Test-Stall-Seconds" in headers:
+        duration = float(headers["X-Datadog-Test-Stall-Seconds"])
+    else:
+        duration = request.app["trace_request_delay"]
+    if duration > 0:
+        log.info("Stalling for %r seconds.", duration)
+        await asyncio.sleep(duration)
+    return v2_decode(data)
+
+
 def v2_decode(data: bytes) -> TelemetryEvent:
     """Decode v2 apm telemetry request data as a dict"""
     # TODO: Handle decoding into a telemetry payload object

releasenotes/notes/Settings-endpoint-928ae5b6e8ddc625.yaml

@@ -0,0 +1,7 @@
+---
+features:
+  - |
+    Added a /test/settings endpoint to change the settings on the fly.
+    It takes a json encoded list of key/values, that will be applied to the app dict.
+    Example:
+    { "trace_request_delay": 5 }

releasenotes/notes/Telemetry-requests-delay-2586180b286ce393.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    --trace-request-delay now also affects telemetry requests.
+    The delay can be set independently for each request by setting the X-Datadog-Test-Stall-Seconds header.

tests/test_agent.py

@@ -400,3 +400,35 @@ def test_uds(tmp_path, available_port):
         pass
     else:
         raise Exception("Test agent failed to start")
+
+
+async def test_post_known_settings(agent):
+    resp = await agent.post(
+        "/test/settings",
+        data='{ "trace_request_delay": 5 }',
+    )
+
+    assert resp.status == 202, await resp.text()
+    assert agent.app["trace_request_delay"] == 5
+
+    resp = await agent.post(
+        "/test/settings",
+        data='{ "trace_request_delay": 0 }',
+    )
+
+    assert resp.status == 202, await resp.text()
+    assert agent.app["trace_request_delay"] == 0
+
+
+async def test_post_unknown_settings(
+    agent,
+):
+    resp = await agent.post(
+        "/test/settings",
+        data='{ "dummy_setting": 5 }',
+    )
+
+    assert resp.status == 422
+    text = await resp.text()
+    assert text == "Unknown key: 'dummy_setting'"
+    assert "dummy_setting" not in agent.app