Skip to content

Commit

Permalink
(chore): add error string to handle_snapshot
Browse files Browse the repository at this point in the history
Prior to this change debugging the handle_snapshot function was
difficult, because no error message beyond "INTERNAL SERVER ERROR" was
returned.  After this change, the exception's message is returned.
  • Loading branch information
tabgok committed Feb 12, 2024
1 parent 97d78ff commit e5b5dad
Showing 1 changed file with 96 additions and 93 deletions.
189 changes: 96 additions & 93 deletions ddapm_test_agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -699,105 +699,108 @@ async def handle_session_start(self, request: Request) -> web.Response:

async def handle_snapshot(self, request: Request) -> web.Response:
"""Generate a snapshot or perform a snapshot test."""
token = request["session_token"]
snap_dir = request.url.query.get("dir", request.app["snapshot_dir"])
snap_ci_mode = request.app["snapshot_ci_mode"]
log.info(
"performing snapshot with token=%r, ci_mode=%r and snapshot directory=%r",
token,
snap_ci_mode,
snap_dir,
)
try:
token = request["session_token"]
snap_dir = request.url.query.get("dir", request.app["snapshot_dir"])
snap_ci_mode = request.app["snapshot_ci_mode"]
log.info(
"performing snapshot with token=%r, ci_mode=%r and snapshot directory=%r",
token,
snap_ci_mode,
snap_dir,
)

# Get the span attributes that are to be ignored for this snapshot.
default_span_ignores: Set[str] = request.app["snapshot_ignored_attrs"]
overrides = set(_parse_csv(request.url.query.get("ignores", "")))
span_ignores = list(default_span_ignores | overrides)
log.info("using ignores %r", span_ignores)

# Get the span attributes that are to be removed for this snapshot.
default_span_removes: Set[str] = request.app["snapshot_removed_attrs"]
overrides = set(_parse_csv(request.url.query.get("removes", "")))
span_removes = list(default_span_removes | overrides)
log.info("using removes %r", span_removes)

if "span_id" in span_removes:
raise AssertionError("Cannot remove 'span_id' from spans")

with CheckTrace.add_frame(f"snapshot (token='{token}')") as frame:
frame.add_item(f"Directory: {snap_dir}")
frame.add_item(f"CI mode: {snap_ci_mode}")

if "X-Datadog-Test-Snapshot-Filename" in request.headers:
snap_file = request.headers["X-Datadog-Test-Snapshot-Filename"]
elif "file" in request.url.query:
snap_file = request.url.query["file"]
else:
snap_file = os.path.join(snap_dir, token)
# Get the span attributes that are to be ignored for this snapshot.
default_span_ignores: Set[str] = request.app["snapshot_ignored_attrs"]
overrides = set(_parse_csv(request.url.query.get("ignores", "")))
span_ignores = list(default_span_ignores | overrides)
log.info("using ignores %r", span_ignores)

# Get the span attributes that are to be removed for this snapshot.
default_span_removes: Set[str] = request.app["snapshot_removed_attrs"]
overrides = set(_parse_csv(request.url.query.get("removes", "")))
span_removes = list(default_span_removes | overrides)
log.info("using removes %r", span_removes)

if "span_id" in span_removes:
raise AssertionError("Cannot remove 'span_id' from spans")

with CheckTrace.add_frame(f"snapshot (token='{token}')") as frame:
frame.add_item(f"Directory: {snap_dir}")
frame.add_item(f"CI mode: {snap_ci_mode}")

if "X-Datadog-Test-Snapshot-Filename" in request.headers:
snap_file = request.headers["X-Datadog-Test-Snapshot-Filename"]
elif "file" in request.url.query:
snap_file = request.url.query["file"]
else:
snap_file = os.path.join(snap_dir, token)

# The logic from here is mostly duplicated for traces and trace stats.
# If another data type is to be snapshotted then it probably makes sense to abstract away
# the required pieces of snapshotting (loading, generating and comparing).
# The logic from here is mostly duplicated for traces and trace stats.
# If another data type is to be snapshotted then it probably makes sense to abstract away
# the required pieces of snapshotting (loading, generating and comparing).

# For backwards compatibility traces don't have a postfix of `_trace.json`
trace_snap_file = f"{snap_file}.json"
tracestats_snap_file = f"{snap_file}_tracestats.json"
# For backwards compatibility traces don't have a postfix of `_trace.json`
trace_snap_file = f"{snap_file}.json"
tracestats_snap_file = f"{snap_file}_tracestats.json"

frame.add_item(f"Trace File: {trace_snap_file}")
frame.add_item(f"Stats File: {tracestats_snap_file}")
log.info("using snapshot files %r and %r", trace_snap_file, tracestats_snap_file)
frame.add_item(f"Trace File: {trace_snap_file}")
frame.add_item(f"Stats File: {tracestats_snap_file}")
log.info("using snapshot files %r and %r", trace_snap_file, tracestats_snap_file)

trace_snap_path_exists = os.path.exists(trace_snap_file)
trace_snap_path_exists = os.path.exists(trace_snap_file)

received_traces = await self._traces_by_session(token)
if snap_ci_mode and received_traces and not trace_snap_path_exists:
raise AssertionError(
f"Trace snapshot file '{trace_snap_file}' not found. "
"Perhaps the file was not checked into source control? "
"The snapshot file is automatically generated when the test agent is not in CI mode."
)
elif trace_snap_path_exists:
# Do the snapshot comparison
with open(trace_snap_file, mode="r") as f:
raw_snapshot = json.load(f)
trace_snapshot.snapshot(
expected_traces=raw_snapshot,
received_traces=received_traces,
ignored=span_ignores,
)
elif received_traces:
# Create a new snapshot for the data received
with open(trace_snap_file, mode="w") as f:
f.write(trace_snapshot.generate_snapshot(received_traces=received_traces, removed=span_removes))
log.info("wrote new trace snapshot to %r", os.path.abspath(trace_snap_file))

# Get all stats buckets from the payloads since we don't care about the other fields (hostname, env, etc)
# in the payload.
received_stats = [bucket for p in (await self._tracestats_by_session(token)) for bucket in p["Stats"]]
tracestats_snap_path_exists = os.path.exists(tracestats_snap_file)
if snap_ci_mode and received_stats and not tracestats_snap_path_exists:
raise AssertionError(
f"Trace stats snapshot file '{tracestats_snap_file}' not found. "
"Perhaps the file was not checked into source control? "
"The snapshot file is automatically generated when the test case is run when not in CI mode."
)
elif tracestats_snap_path_exists:
# Do the snapshot comparison
with open(tracestats_snap_file, mode="r") as f:
raw_snapshot = json.load(f)
tracestats_snapshot.snapshot(
expected_stats=raw_snapshot,
received_stats=received_stats,
)
elif received_stats:
# Create a new snapshot for the data received
with open(tracestats_snap_file, mode="w") as f:
f.write(tracestats_snapshot.generate(received_stats))
log.info(
"wrote new tracestats snapshot to %r",
os.path.abspath(tracestats_snap_file),
)
return web.HTTPOk()
received_traces = await self._traces_by_session(token)
if snap_ci_mode and received_traces and not trace_snap_path_exists:
raise AssertionError(
f"Trace snapshot file '{trace_snap_file}' not found. "
"Perhaps the file was not checked into source control? "
"The snapshot file is automatically generated when the test agent is not in CI mode."
)
elif trace_snap_path_exists:
# Do the snapshot comparison
with open(trace_snap_file, mode="r") as f:
raw_snapshot = json.load(f)
trace_snapshot.snapshot(
expected_traces=raw_snapshot,
received_traces=received_traces,
ignored=span_ignores,
)
elif received_traces:
# Create a new snapshot for the data received
with open(trace_snap_file, mode="w") as f:
f.write(trace_snapshot.generate_snapshot(received_traces=received_traces, removed=span_removes))
log.info("wrote new trace snapshot to %r", os.path.abspath(trace_snap_file))

# Get all stats buckets from the payloads since we don't care about the other fields (hostname, env, etc)
# in the payload.
received_stats = [bucket for p in (await self._tracestats_by_session(token)) for bucket in p["Stats"]]
tracestats_snap_path_exists = os.path.exists(tracestats_snap_file)
if snap_ci_mode and received_stats and not tracestats_snap_path_exists:
raise AssertionError(
f"Trace stats snapshot file '{tracestats_snap_file}' not found. "
"Perhaps the file was not checked into source control? "
"The snapshot file is automatically generated when the test case is run when not in CI mode."
)
elif tracestats_snap_path_exists:
# Do the snapshot comparison
with open(tracestats_snap_file, mode="r") as f:
raw_snapshot = json.load(f)
tracestats_snapshot.snapshot(
expected_stats=raw_snapshot,
received_stats=received_stats,
)
elif received_stats:
# Create a new snapshot for the data received
with open(tracestats_snap_file, mode="w") as f:
f.write(tracestats_snapshot.generate(received_stats))
log.info(
"wrote new tracestats snapshot to %r",
os.path.abspath(tracestats_snap_file),
)
return web.HTTPOk()
except Exception as e:
return web.HTTPInternalServerError(reason=str(e))

async def handle_session_traces(self, request: Request) -> web.Response:
token = request["session_token"]
Expand Down

0 comments on commit e5b5dad

Please sign in to comment.