Skip to content

Commit

Permalink
Unit tests for feature enable/disable and refactor logic
Browse files Browse the repository at this point in the history 
Signed-off-by: Siddhant Deshmukh <[email protected]>
  • Loading branch information
@deshsidd
deshsidd committed Jul 17, 2024
1 parent d2ba20b commit fb4ce18
Show file tree
Hide file tree
Showing 4 changed files with 100 additions and 23 deletions.
15 changes: 6 additions & 9 deletions src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Expand Down
6 changes: 3 additions & 3 deletions src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Expand Down
76 changes: 65 additions & 11 deletions src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
}
}
26 changes: 26 additions & 0 deletions src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}.
Expand All @@ -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() {
Expand All @@ -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() {
Expand Down Expand Up @@ -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();
}
}

0 comments on commit fb4ce18

Please sign in to comment.