pass new argument in both possible cases

ddapm_test_agent/agent.py

@@ -750,6 +750,7 @@ async def handle_snapshot(self, request: Request) -> web.Response:
             trace_snap_path_exists = os.path.exists(trace_snap_file)
 
             received_traces = await self._traces_by_session(token)
+            count_tolerance = int(request.query.get("count_tolerance", 0))
             if snap_ci_mode and received_traces and not trace_snap_path_exists:
                 raise AssertionError(
                     f"Trace snapshot file '{trace_snap_file}' not found. "
@@ -764,6 +765,7 @@ async def handle_snapshot(self, request: Request) -> web.Response:
                     expected_traces=raw_snapshot,
                     received_traces=received_traces,
                     ignored=span_ignores,
+                    count_tolerance=count_tolerance,
                 )
             elif received_traces:
                 # Create a new snapshot for the data received
@@ -786,9 +788,7 @@ async def handle_snapshot(self, request: Request) -> web.Response:
                 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,
-                    count_tolerance=int(request.query.get("count_tolerance", 0)),
+                    expected_stats=raw_snapshot, received_stats=received_stats, count_tolerance=count_tolerance
                 )
             elif received_stats:
                 # Create a new snapshot for the data received