Update unit test

apache · Jan 23, 2025 · 562c21f · 562c21f
1 parent f6c40fb
commit 562c21f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 10 deletions.
diff --git a/.../src/main/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainer.java b/.../src/main/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainer.java
@@ -34,6 +34,7 @@
 import org.apache.beam.runners.core.metrics.GaugeCell;
 import org.apache.beam.runners.core.metrics.MetricsMap;
 import org.apache.beam.runners.core.metrics.StringSetCell;
+import org.apache.beam.runners.core.metrics.StringSetData;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Distribution;
 import org.apache.beam.sdk.metrics.Gauge;
@@ -69,7 +70,7 @@ public class StreamingStepMetricsContainer implements MetricsContainer {
 
   private MetricsMap<MetricName, GaugeCell> gauges = new MetricsMap<>(GaugeCell::new);
 
-  private MetricsMap<MetricName, StringSetCell> stringSet = new MetricsMap<>(StringSetCell::new);
+  private MetricsMap<MetricName, StringSetCell> stringSets = new MetricsMap<>(StringSetCell::new);
 
   private MetricsMap<MetricName, DeltaDistributionCell> distributions =
       new MetricsMap<>(DeltaDistributionCell::new);
@@ -165,7 +166,7 @@ public Gauge getGauge(MetricName metricName) {
 
   @Override
   public StringSet getStringSet(MetricName metricName) {
-    return stringSet.get(metricName);
+    return stringSets.get(metricName);
   }
 
   @Override
@@ -188,7 +189,8 @@ public Iterable<CounterUpdate> extractUpdates() {
     // Streaming metrics are updated as delta and not cumulative.
     return counterUpdates()
         .append(distributionUpdates())
-        .append(gaugeUpdates().append(stringSetUpdates()));
+        .append(gaugeUpdates())
+        .append(stringSetUpdates());
   }
 
   private FluentIterable<CounterUpdate> counterUpdates() {
@@ -231,15 +233,18 @@ private FluentIterable<CounterUpdate> gaugeUpdates() {
   }
 
   private FluentIterable<CounterUpdate> stringSetUpdates() {
-    return FluentIterable.from(stringSet.entries())
+    return FluentIterable.from(stringSets.entries())
         .transform(
             new Function<Entry<MetricName, StringSetCell>, CounterUpdate>() {
               @Override
               public @Nullable CounterUpdate apply(
                   @Nonnull Map.Entry<MetricName, StringSetCell> entry) {
+                StringSetData value = entry.getValue().getAndReset();
+                if (value.stringSet().isEmpty()) {
+                  return null;
+                }
                 return MetricsToCounterUpdateConverter.fromStringSet(
-                    MetricKey.create(stepName, entry.getKey()),
-                    false, entry.getValue().getAndReset());
+                    MetricKey.create(stepName, entry.getKey()), false, value);
               }
             })
         .filter(Predicates.notNull());

diff --git a/.../test/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainerTest.java b/.../test/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainerTest.java
@@ -315,14 +315,14 @@ public void testStringSetUpdateExtraction() {
             .setStringList(new StringList().setElements(Arrays.asList("ij", "kl", "mn")));
 
     updates = StreamingStepMetricsContainer.extractMetricUpdates(registry);
-    assertThat(updates, containsInAnyOrder(name1Update, name2Update));
+    assertThat(updates, containsInAnyOrder(name2Update));
 
+    // test deltas
     c1.getStringSet(name1).add("op");
-    name1Update.setStringList(
-        new StringList().setElements(Arrays.asList("ab", "cd", "ef", "gh", "op")));
+    name1Update.setStringList(new StringList().setElements(Arrays.asList("op")));
 
     updates = StreamingStepMetricsContainer.extractMetricUpdates(registry);
-    assertThat(updates, containsInAnyOrder(name1Update, name2Update));
+    assertThat(updates, containsInAnyOrder(name1Update));
   }
 
   @Test