From fb4ce18cfad8befe1795bbc90b10c67537cdcc1d Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Tue, 16 Jul 2024 23:49:18 -0700 Subject: [PATCH] Unit tests for feature enable/disable and refactor logic Signed-off-by: Siddhant Deshmukh --- .../core/listener/QueryInsightsListener.java | 15 ++-- .../core/service/QueryInsightsService.java | 6 +- .../listener/QueryInsightsListenerTests.java | 76 ++++++++++++++++--- .../service/QueryInsightsServiceTests.java | 26 +++++++ 4 files changed, 100 insertions(+), 23 deletions(-) diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java index 0e5cd7a..c54f1ad 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java +++ b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java @@ -94,21 +94,18 @@ public QueryInsightsListener(final ClusterService clusterService, final QueryIns * @param isCurrentMetricEnabled boolean */ public void setEnableTopQueries(final MetricType metricType, final boolean isCurrentMetricEnabled) { - boolean isTopNFeatureDisabled = !queryInsightsService.isTopNFeatureEnabled(); + boolean isTopNFeaturePreviouslyDisabled = !queryInsightsService.isTopNFeatureEnabled(); this.queryInsightsService.enableCollection(metricType, isCurrentMetricEnabled); + boolean isTopNFeatureCurrentlyDisabled = !queryInsightsService.isTopNFeatureEnabled(); - if (!isCurrentMetricEnabled) { - // disable QueryInsightsListener only if all metrics collections are disabled now - // and search query metrics is disabled. - if (isTopNFeatureDisabled) { - super.setEnabled(false); + if (isTopNFeatureCurrentlyDisabled) { + super.setEnabled(false); + if (!isTopNFeaturePreviouslyDisabled) { queryInsightsService.checkAndStopQueryInsights(); } } else { super.setEnabled(true); - // restart QueryInsightsListener only if none of metrics collections is enabled before and - // search query metrics is disabled before. - if (isTopNFeatureDisabled) { + if (isTopNFeaturePreviouslyDisabled) { queryInsightsService.checkAndRestartQueryInsights(); } } diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java index 4c24736..6a36c9a 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java @@ -310,14 +310,14 @@ public void setExporter(final MetricType type, final Settings settings) { * @param searchQueryMetricsEnabled boolean flag */ public void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) { - boolean oldsearchQueryMetricsEnabled = isSearchQueryMetricsFeatureEnabled(); + boolean oldSearchQueryMetricsEnabled = isSearchQueryMetricsFeatureEnabled(); this.searchQueryMetricsEnabled = searchQueryMetricsEnabled; if (searchQueryMetricsEnabled) { - if (!oldsearchQueryMetricsEnabled) { + if (!oldSearchQueryMetricsEnabled) { checkAndRestartQueryInsights(); } } else { - if (oldsearchQueryMetricsEnabled) { + if (oldSearchQueryMetricsEnabled) { checkAndStopQueryInsights(); } } diff --git a/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java b/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java index 93876c6..6db7ffe 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java @@ -10,12 +10,15 @@ import org.junit.Before; import org.mockito.ArgumentCaptor; +import org.mockito.MockitoAnnotations; +import org.mockito.MockitoSession; import org.opensearch.action.search.SearchPhaseContext; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchRequestContext; import org.opensearch.action.search.SearchTask; import org.opensearch.action.search.SearchType; import org.opensearch.action.support.replication.ClusterStateCreationUtils; +import org.opensearch.client.Client; import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.collect.Tuple; @@ -34,6 +37,8 @@ import org.opensearch.search.aggregations.support.ValueType; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.tasks.Task; +import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.metrics.tags.Tags; import org.opensearch.test.ClusterServiceUtils; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; @@ -50,7 +55,10 @@ import java.util.concurrent.TimeUnit; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -201,16 +209,62 @@ public void testConcurrentOnRequestEnd() throws InterruptedException { verify(queryInsightsService, times(numRequests)).addRecord(any()); } - public void testSetEnabled() { - when(queryInsightsService.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true); - QueryInsightsListener queryInsightsListener = new QueryInsightsListener(clusterService, queryInsightsService); - queryInsightsListener.setEnableTopQueries(MetricType.LATENCY, true); - assertTrue(queryInsightsListener.isEnabled()); - - when(queryInsightsService.isCollectionEnabled(MetricType.LATENCY)).thenReturn(false); - when(queryInsightsService.isCollectionEnabled(MetricType.CPU)).thenReturn(false); - when(queryInsightsService.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); - queryInsightsListener.setEnableTopQueries(MetricType.LATENCY, false); - assertFalse(queryInsightsListener.isEnabled()); + public void testTopNFeatureEnabledDisabled() { + // Test case 1: Only latency enabled initially, disable latency and verify expected behavior + QueryInsightsService queryInsightsService1 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener1 = new QueryInsightsListener(clusterService, queryInsightsService1); + + when(queryInsightsService1.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true); + when(queryInsightsService1.isCollectionEnabled(MetricType.CPU)).thenReturn(false); + when(queryInsightsService1.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService1.isTopNFeatureEnabled()).thenReturn(true).thenReturn(false); + + queryInsightsListener1.setEnableTopQueries(MetricType.LATENCY, false); + assertFalse(queryInsightsListener1.isEnabled()); + verify(queryInsightsService1).checkAndStopQueryInsights(); + verify(queryInsightsService1, never()).checkAndRestartQueryInsights(); + + // Test case 2: All disabled initially, enable latency and verify expected behavior + QueryInsightsService queryInsightsService2 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener2 = new QueryInsightsListener(clusterService, queryInsightsService2); + + when(queryInsightsService2.isCollectionEnabled(MetricType.LATENCY)).thenReturn(false); + when(queryInsightsService2.isCollectionEnabled(MetricType.CPU)).thenReturn(false); + when(queryInsightsService2.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService2.isTopNFeatureEnabled()).thenReturn(false).thenReturn(true); + + queryInsightsListener2.setEnableTopQueries(MetricType.LATENCY, true); + assertTrue(queryInsightsListener2.isEnabled()); + verify(queryInsightsService2).checkAndRestartQueryInsights(); + verify(queryInsightsService2, never()).checkAndStopQueryInsights(); + + // Test case 3: latency and CPU enabled initially, disable latency and verify expected behavior + QueryInsightsService queryInsightsService3 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener3 = new QueryInsightsListener(clusterService, queryInsightsService3); + + when(queryInsightsService3.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true); + when(queryInsightsService3.isCollectionEnabled(MetricType.CPU)).thenReturn(true); + when(queryInsightsService3.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService3.isTopNFeatureEnabled()).thenReturn(true).thenReturn(true); + + queryInsightsListener3.setEnableTopQueries(MetricType.LATENCY, false); + assertTrue(queryInsightsListener3.isEnabled()); + verify(queryInsightsService3, never()).checkAndStopQueryInsights(); + verify(queryInsightsService3, never()).checkAndRestartQueryInsights(); + + + // Test case 4: Only CPU enabled initially, enable latency and verify expected behavior + QueryInsightsService queryInsightsService4 = mock(QueryInsightsService.class); + QueryInsightsListener queryInsightsListener4 = new QueryInsightsListener(clusterService, queryInsightsService4); + + when(queryInsightsService4.isCollectionEnabled(MetricType.LATENCY)).thenReturn(false); + when(queryInsightsService4.isCollectionEnabled(MetricType.CPU)).thenReturn(true); + when(queryInsightsService4.isCollectionEnabled(MetricType.MEMORY)).thenReturn(false); + when(queryInsightsService4.isTopNFeatureEnabled()).thenReturn(true).thenReturn(true); + + queryInsightsListener4.setEnableTopQueries(MetricType.LATENCY, true); + assertTrue(queryInsightsListener4.isEnabled()); + verify(queryInsightsService4, never()).checkAndRestartQueryInsights(); + verify(queryInsightsService4, never()).checkAndStopQueryInsights(); } } diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java index bf273ad..cd7a7a7 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java @@ -21,6 +21,10 @@ import org.opensearch.threadpool.ThreadPool; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; /** * Unit Tests for {@link QueryInsightsService}. @@ -29,6 +33,7 @@ public class QueryInsightsServiceTests extends OpenSearchTestCase { private final ThreadPool threadPool = mock(ThreadPool.class); private final Client client = mock(Client.class); private QueryInsightsService queryInsightsService; + private QueryInsightsService queryInsightsServiceSpy; @Before public void setup() { @@ -40,6 +45,7 @@ public void setup() { queryInsightsService.enableCollection(MetricType.LATENCY, true); queryInsightsService.enableCollection(MetricType.CPU, true); queryInsightsService.enableCollection(MetricType.MEMORY, true); + queryInsightsServiceSpy = spy(queryInsightsService); } public void testAddRecordToLimitAndDrain() { @@ -84,4 +90,24 @@ public void testSearchQueryMetricsEnabled() { assertNotNull(queryInsightsService.getSearchQueryCategorizer()); } + + public void testFeaturesEnableDisable() { + // Test case 1: All metric type collection disabled and search query metrics disabled, enable search query metrics + queryInsightsServiceSpy.enableCollection(MetricType.LATENCY, false); + queryInsightsServiceSpy.enableCollection(MetricType.CPU, false); + queryInsightsServiceSpy.enableCollection(MetricType.MEMORY, false); + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(false); + + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(true); + verify(queryInsightsServiceSpy).checkAndRestartQueryInsights(); + + // Test case 2: All metric type collection disabled and search query metrics enabled, disable search query metrics + queryInsightsServiceSpy.enableCollection(MetricType.LATENCY, false); + queryInsightsServiceSpy.enableCollection(MetricType.CPU, false); + queryInsightsServiceSpy.enableCollection(MetricType.MEMORY, false); + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(true); + + queryInsightsServiceSpy.setSearchQueryMetricsEnabled(false); + verify(queryInsightsServiceSpy).checkAndStopQueryInsights(); + } }