Add /test/session/tracerflares endpoint to return all tracer-flares received by the test agent for a given session token. (#146)

... received by the test agent for a given session token.

Author: mcculls

README.md

@@ -408,6 +408,31 @@ Mimics the tracer_flare endpoint of the agent. Returns OK if the flare contains
 
 Logs a line everytime it's called and stores the tracer flare details in the request under `"_tracer_flare"`.
 
+### /test/session/tracerflares
+
+Return all tracer-flares that have been received by the agent for the given session token.
+
+#### [optional] `?test_session_token=`
+#### [optional] `X-Datadog-Test-Session-Token`
+
+Returns the tracer-flares in the following json format:
+
+```json
+[
+  {
+    "source": "...",
+    "case_id": "...",
+    "email": "...",
+    "hostname": "...",
+    "flare_file": "...",
+  }
+]
+```
+
+`flare_file` is the base64 encoded content of the tracer-flare payload.
+
+If there was an error parsing the tracer-flare form, that will be recorded under `error`.
+
 ## Development
 
 ### Prerequisites

ddapm_test_agent/agent.py

@@ -1,5 +1,4 @@
 import argparse
-from asyncio import StreamReader
 import atexit
 import base64
 from collections import OrderedDict
@@ -10,7 +9,6 @@
 import pprint
 import socket
 import sys
-from typing import Any
 from typing import Awaitable
 from typing import Callable
 from typing import Dict
@@ -24,7 +22,6 @@
 from urllib.parse import urlunparse
 
 from aiohttp import ClientSession
-from aiohttp import MultipartReader
 from aiohttp import web
 from aiohttp.web import Request
 from aiohttp.web import middleware
@@ -54,6 +51,8 @@
 from .trace_checks import CheckTraceDDService
 from .trace_checks import CheckTracePeerService
 from .trace_checks import CheckTraceStallAsync
+from .tracerflare import TracerFlareEvent
+from .tracerflare import v1_decode as v1_tracerflare_decode
 from .tracestats import decode_v06 as tracestats_decode_v06
 from .tracestats import v06StatsPayload
 
@@ -381,6 +380,19 @@ async def _apmtelemetry_by_session(self, token: Optional[str]) -> List[Telemetry
         # TODO: Sort the events?
         return events
 
+    async def _tracerflares_by_session(self, token: Optional[str]) -> List[TracerFlareEvent]:
+        """Return the tracer-flare events that belong to the given session token.
+
+        If token is None or if the token was used to manually start a session
+        with /session-start then return all tracer-flare events that were sent
+        since the last /session-start request was made.
+        """
+        events: List[TracerFlareEvent] = []
+        for req in self._requests_by_session(token):
+            if req.match_info.handler == self.handle_v1_tracer_flare:
+                events.append(await v1_tracerflare_decode(req.headers, await req.read()))
+        return events
+
     async def _tracestats_by_session(self, token: Optional[str]) -> List[v06StatsPayload]:
         stats: List[v06StatsPayload] = []
         for req in self._requests_by_session(token):
@@ -517,21 +529,7 @@ async def handle_v2_apmtelemetry(self, request: Request) -> web.Response:
         return web.HTTPOk()
 
     async def handle_v1_tracer_flare(self, request: Request) -> web.Response:
-        # reconstruct stream from previously cached bytes
-        stream = StreamReader()
-        stream.feed_data(self._request_data(request))
-        stream.feed_eof()
-
-        tracer_flare: Dict[str, Any] = {}
-
-        async for part in MultipartReader(request.headers, stream):
-            if part.name is not None:
-                if part.name == "flare_file":
-                    tracer_flare[part.name] = await part.read()  # zipfile
-                else:
-                    tracer_flare[part.name] = await part.text()
-
-        request["_tracer_flare"] = tracer_flare
+        tracer_flare: TracerFlareEvent = await v1_tracerflare_decode(request.headers, self._request_data(request))
 
         expectedFields = ["source", "case_id", "email", "hostname", "flare_file"]
         missingFields = [k for k in expectedFields if k not in tracer_flare]
@@ -782,6 +780,11 @@ async def handle_session_apmtelemetry(self, request: Request) -> web.Response:
         events = await self._apmtelemetry_by_session(token)
         return web.json_response(events)
 
+    async def handle_session_tracerflares(self, request: Request) -> web.Response:
+        token = request["session_token"]
+        events = await self._tracerflares_by_session(token)
+        return web.json_response(events)
+
     async def handle_session_tracestats(self, request: Request) -> web.Response:
         token = request["session_token"]
         stats = await self._tracestats_by_session(token)
@@ -1029,6 +1032,7 @@ def make_app(
             web.get("/test/session/snapshot", agent.handle_snapshot),
             web.get("/test/session/traces", agent.handle_session_traces),
             web.get("/test/session/apmtelemetry", agent.handle_session_apmtelemetry),
+            web.get("/test/session/tracerflares", agent.handle_session_tracerflares),
             web.get("/test/session/stats", agent.handle_session_tracestats),
             web.get("/test/session/requests", agent.handle_session_requests),
             web.post("/test/session/responses/config", agent.handle_v07_remoteconfig_create),

ddapm_test_agent/tracerflare.py

@@ -0,0 +1,27 @@
+from asyncio import StreamReader
+import base64
+from typing import Dict
+from typing import Mapping
+
+from aiohttp import MultipartReader
+
+
+TracerFlareEvent = Dict[str, str]
+
+
+async def v1_decode(headers: Mapping[str, str], data: bytes) -> TracerFlareEvent:
+    """Decode v1 tracer flare form as a dict"""
+    tracer_flare: TracerFlareEvent = {}
+    try:
+        stream = StreamReader()
+        stream.feed_data(data)
+        stream.feed_eof()
+        async for part in MultipartReader(headers, stream):
+            if part.name is not None:
+                if part.name == "flare_file":
+                    tracer_flare[part.name] = base64.b64encode(await part.read()).decode("ascii")
+                else:
+                    tracer_flare[part.name] = await part.text()
+    except Exception as err:
+        tracer_flare["error"] = str(err)
+    return tracer_flare

releasenotes/notes/session-tracerflares-400ce9b3d3838b40.yaml

@@ -0,0 +1,4 @@
+features:
+  - |
+    Add ``/test/session/tracerflares`` endpoint to return all tracer-flares received
+    by the test agent for a given session token.

tests/test_tracerflare.py

@@ -2,6 +2,13 @@
 
 
 async def test_tracerflare(agent):
+    expected_output = {
+        "source": "tracer_test",
+        "case_id": "12345",
+        "email": "[email protected]",
+        "hostname": "my.hostname",
+        "flare_file": "UEsFBgAAAAAAAAAAAAAAAAAAAAAAAA==",
+    }
     form = FormData()
     form.add_field("source", "tracer_test")
     form.add_field("case_id", "12345")
@@ -15,9 +22,17 @@ async def test_tracerflare(agent):
     )
     resp = await agent.post("/tracer_flare/v1", data=form)
     assert resp.status == 200
+    flares = await agent.get("/test/session/tracerflares")
+    assert await flares.json() == [expected_output]
 
 
 async def test_tracerflare_missing_case_id(agent):
+    expected_output = {
+        "source": "tracer_test",
+        "email": "[email protected]",
+        "hostname": "my.hostname",
+        "flare_file": "UEsFBgAAAAAAAAAAAAAAAAAAAAAAAA==",
+    }
     form = FormData()
     form.add_field("source", "tracer_test")
     form.add_field("email", "[email protected]")
@@ -30,13 +45,21 @@ async def test_tracerflare_missing_case_id(agent):
     )
     resp = await agent.post("/tracer_flare/v1", data=form)
     assert resp.status == 400
+    flares = await agent.get("/test/session/tracerflares")
+    assert await flares.json() == [expected_output]
 
 
-async def test_tracerflare_missing_file(agent):
-    form = FormData()
-    form.add_field("source", "tracer_test")
-    form.add_field("case_id", "12345")
-    form.add_field("email", "[email protected]")
-    form.add_field("hostname", "my.hostname")
+async def test_tracerflare_not_multipart(agent):
+    expected_output = {
+        "error": "multipart/* content type expected",
+    }
+    form = {
+        "source": "tracer_test",
+        "case_id": "12345",
+        "email": "[email protected]",
+        "hostname": "my.hostname",
+    }
     resp = await agent.post("/tracer_flare/v1", data=form)
     assert resp.status == 400
+    flares = await agent.get("/test/session/tracerflares")
+    assert await flares.json() == [expected_output]