we just need to store the counts

evanchooly · evanchooly · commit 0df836a77171 · 2025-01-24T18:51:32.000-05:00
diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/LogProbeTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/LogProbeTest.java
@@ -29,8 +29,6 @@
 import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.core.CoreTracer;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -84,23 +82,24 @@ public void noDebugSession() {
   @Test
   public void budgets() {
     BudgetSink sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
-
     DebuggerAgentHelper.injectSink(sink);
-    assertEquals(0, sink.captures.size());
+
     TracerAPI tracer =
         CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build();
     AgentTracer.registerIfAbsent(tracer);
     int runs = 100;
     for (int i = 0; i < runs; i++) {
       runTrace(tracer, true);
     }
-    assertEquals(runs * LogProbe.CAPTURING_PROBE_BUDGET, sink.captures.size());
+    assertEquals(runs * LogProbe.CAPTURING_PROBE_BUDGET, sink.captures);
 
-    runs = 2000;
+    sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
+    DebuggerAgentHelper.injectSink(sink);
+    runs = 1010;
     for (int i = 0; i < runs; i++) {
       runTrace(tracer, false);
     }
-    assertEquals(runs * LogProbe.NON_CAPTURING_PROBE_BUDGET, sink.highRate.size());
+    assertEquals(runs * LogProbe.NON_CAPTURING_PROBE_BUDGET, sink.highRate);
   }
 
   private void runTrace(TracerAPI tracer, boolean captureSnapshot) {
@@ -312,21 +311,21 @@ private Builder createLog(String template) {
 
   private static class BudgetSink extends DebuggerSink {
 
-    public List<Snapshot> captures = new ArrayList<>();
-    public List<Snapshot> highRate = new ArrayList<>();
+    public int captures;
+    public int highRate;
 
     public BudgetSink(Config config, ProbeStatusSink probeStatusSink) {
       super(config, probeStatusSink);
     }
 
     @Override
     public void addHighRateSnapshot(Snapshot snapshot) {
-      highRate.add(snapshot);
+      highRate++;
     }
 
     @Override
     public void addSnapshot(Snapshot snapshot) {
-      captures.add(snapshot);
+      captures++;
     }
 
     @Override
