Add tests for checking attributes in metrics

carlosdelest · Sep 16, 2024 · 9b0af4d · 9b0af4d
1 parent 3dc18de
commit 9b0af4d
Showing 1 changed file with 46 additions and 8 deletions.
diff --git a/...rc/test/java/org/elasticsearch/search/TelemetryMetrics/SearchTransportTelemetryTests.java b/...rc/test/java/org/elasticsearch/search/TelemetryMetrics/SearchTransportTelemetryTests.java
@@ -24,13 +24,16 @@
 import java.util.Collection;
 import java.util.List;
 
+import static org.elasticsearch.action.search.SearchTransportAPMMetrics.ACTION_ATTRIBUTE_NAME;
 import static org.elasticsearch.action.search.SearchTransportAPMMetrics.DFS_ACTION_METRIC;
 import static org.elasticsearch.action.search.SearchTransportAPMMetrics.FETCH_ID_ACTION_METRIC;
 import static org.elasticsearch.action.search.SearchTransportAPMMetrics.FETCH_ID_SCROLL_ACTION_METRIC;
 import static org.elasticsearch.action.search.SearchTransportAPMMetrics.FREE_CONTEXT_SCROLL_ACTION_METRIC;
 import static org.elasticsearch.action.search.SearchTransportAPMMetrics.QUERY_ACTION_METRIC;
 import static org.elasticsearch.action.search.SearchTransportAPMMetrics.QUERY_ID_ACTION_METRIC;
 import static org.elasticsearch.action.search.SearchTransportAPMMetrics.QUERY_SCROLL_ACTION_METRIC;
+import static org.elasticsearch.action.search.SearchTransportAPMMetrics.SEARCH_ACTION_LATENCY_BASE_METRIC;
+import static org.elasticsearch.action.search.SearchTransportAPMMetrics.SYSTEM_THREAD_ATTRIBUTE_NAME;
 import static org.elasticsearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE;
 import static org.elasticsearch.index.query.QueryBuilders.simpleQueryStringQuery;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertScrollResponsesAndHitCount;
@@ -121,6 +124,37 @@ public void testSearchTransportMetricsScroll() throws InterruptedException {
         resetMeter();
     }
 
+    public void testSearchTransportMetricsAttributesDfsQueryThenFetch() throws InterruptedException {
+        assertSearchHitsWithoutFailures(
+            client().prepareSearch(indexName).setSearchType(SearchType.DFS_QUERY_THEN_FETCH).setQuery(simpleQueryStringQuery("doc1")),
+            "1"
+        );
+        checkMetricsAttributes();
+    }
+
+    public void testSearchTransportMetricsAttributesQueryThenFetch() throws InterruptedException {
+        assertSearchHitsWithoutFailures(
+            client().prepareSearch(indexName).setSearchType(SearchType.QUERY_THEN_FETCH).setQuery(simpleQueryStringQuery("doc1")),
+            "1"
+        );
+        checkMetricsAttributes();
+    }
+
+    public void testSearchTransportMetricsAttributesScroll() throws InterruptedException {
+        assertScrollResponsesAndHitCount(
+            client(),
+            TimeValue.timeValueSeconds(60),
+            client().prepareSearch(indexName)
+                .setSearchType(SearchType.DFS_QUERY_THEN_FETCH)
+                .setSize(1)
+                .setQuery(simpleQueryStringQuery("doc1 doc2")),
+            2,
+            (a, b) -> {}
+        );
+
+        checkMetricsAttributes();
+    }
+
     private void resetMeter() {
         getTestTelemetryPlugin().resetMeter();
     }
@@ -130,13 +164,17 @@ private TestTelemetryPlugin getTestTelemetryPlugin() {
     }
 
     private long getNumberOfMeasurements(String attributeValue) {
-        final List<Measurement> measurements = getTestTelemetryPlugin().getLongHistogramMeasurement(
-            org.elasticsearch.action.search.SearchTransportAPMMetrics.SEARCH_ACTION_LATENCY_BASE_METRIC
-        );
-        return measurements.stream()
-            .filter(
-                m -> m.attributes().get(org.elasticsearch.action.search.SearchTransportAPMMetrics.ACTION_ATTRIBUTE_NAME) == attributeValue
-            )
-            .count();
+        final List<Measurement> measurements = getTestTelemetryPlugin().getLongHistogramMeasurement(SEARCH_ACTION_LATENCY_BASE_METRIC);
+        return measurements.stream().filter(m -> m.attributes().get(ACTION_ATTRIBUTE_NAME) == attributeValue).count();
+    }
+
+    private void checkMetricsAttributes() {
+        final List<Measurement> measurements = getTestTelemetryPlugin().getLongHistogramMeasurement(SEARCH_ACTION_LATENCY_BASE_METRIC);
+        assertTrue(measurements.stream().allMatch(m -> checkMeasurementAttributes(m)));
+    }
+
+    private boolean checkMeasurementAttributes(Measurement m) {
+        return (m.attributes().get(ACTION_ATTRIBUTE_NAME) != null)
+            && (((boolean) m.attributes().get(SYSTEM_THREAD_ATTRIBUTE_NAME)) == false);
     }
 }