From fe45ecdefb1783c00273128a2f4555050066eb41 Mon Sep 17 00:00:00 2001 From: Manasvini B S Date: Tue, 13 Aug 2024 22:47:35 -0700 Subject: [PATCH] FIx unit tests for PIT changes Signed-off-by: Manasvini B S --- .../unittest/cursor/DefaultCursorTest.java | 58 ++++++++++++++++- .../query/DefaultQueryActionTest.java | 63 +++++++++++++++---- 2 files changed, 109 insertions(+), 12 deletions(-) diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/unittest/cursor/DefaultCursorTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/unittest/cursor/DefaultCursorTest.java index 1b9662035d..8506c7a1e5 100644 --- a/legacy/src/test/java/org/opensearch/sql/legacy/unittest/cursor/DefaultCursorTest.java +++ b/legacy/src/test/java/org/opensearch/sql/legacy/unittest/cursor/DefaultCursorTest.java @@ -9,14 +9,47 @@ import static org.hamcrest.Matchers.emptyOrNullString; import static org.hamcrest.Matchers.startsWith; import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.when; +import java.io.ByteArrayOutputStream; import java.util.ArrayList; import java.util.Collections; +import org.junit.Before; import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.legacy.cursor.CursorType; import org.opensearch.sql.legacy.cursor.DefaultCursor; +import org.opensearch.sql.legacy.esdomain.LocalClusterState; +import org.opensearch.sql.opensearch.setting.OpenSearchSettings; public class DefaultCursorTest { + @Mock private OpenSearchSettings settings; + + @Mock private SearchSourceBuilder sourceBuilder; + + @Before + public void setUp() { + MockitoAnnotations.openMocks(this); + // Required for Pagination queries using PIT instead of Scroll + doReturn(Collections.emptyList()).when(settings).getSettings(); + when(settings.getSettingValue(Settings.Key.SQL_PAGINATION_API_SEARCH_AFTER)).thenReturn(true); + LocalClusterState.state().setPluginSettings(settings); + + // Mock the toXContent method of SearchSourceBuilder + try { + XContentBuilder xContentBuilder = XContentFactory.jsonBuilder(new ByteArrayOutputStream()); + when(sourceBuilder.toXContent(any(XContentBuilder.class), any())).thenReturn(xContentBuilder); + } catch (Exception e) { + throw new RuntimeException(e); + } + } @Test public void checkCursorType() { @@ -25,7 +58,26 @@ public void checkCursorType() { } @Test - public void cursorShouldStartWithCursorTypeID() { + public void cursorShouldStartWithCursorTypeIDForPIT() { + DefaultCursor cursor = new DefaultCursor(); + cursor.setRowsLeft(50); + cursor.setPitId("dbdskbcdjksbcjkdsbcjk+//"); + cursor.setIndexPattern("myIndex"); + cursor.setFetchSize(500); + cursor.setFieldAliasMap(Collections.emptyMap()); + cursor.setColumns(new ArrayList<>()); + + // Set the mocked SearchSourceBuilder to the cursor + cursor.searchSourceBuilder = sourceBuilder; + + assertThat(cursor.generateCursorId(), startsWith(cursor.getType().getId() + ":")); + } + + @Test + public void cursorShouldStartWithCursorTypeIDForScroll() { + // Disable PIT for pagination and use scroll instead + when(settings.getSettingValue(Settings.Key.SQL_PAGINATION_API_SEARCH_AFTER)).thenReturn(false); + DefaultCursor cursor = new DefaultCursor(); cursor.setRowsLeft(50); cursor.setScrollId("dbdskbcdjksbcjkdsbcjk+//"); @@ -33,6 +85,10 @@ public void cursorShouldStartWithCursorTypeID() { cursor.setFetchSize(500); cursor.setFieldAliasMap(Collections.emptyMap()); cursor.setColumns(new ArrayList<>()); + + // Set the mocked SearchSourceBuilder to the cursor + cursor.searchSourceBuilder = sourceBuilder; + assertThat(cursor.generateCursorId(), startsWith(cursor.getType().getId() + ":")); } diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/unittest/query/DefaultQueryActionTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/unittest/query/DefaultQueryActionTest.java index 755d604a65..ec6ab00f97 100644 --- a/legacy/src/test/java/org/opensearch/sql/legacy/unittest/query/DefaultQueryActionTest.java +++ b/legacy/src/test/java/org/opensearch/sql/legacy/unittest/query/DefaultQueryActionTest.java @@ -8,16 +8,9 @@ import static org.hamcrest.Matchers.equalTo; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; - -import java.util.ArrayList; -import java.util.LinkedList; -import java.util.List; -import java.util.Optional; +import static org.mockito.Mockito.*; + +import java.util.*; import org.junit.After; import org.junit.Assert; import org.junit.Before; @@ -27,6 +20,8 @@ import org.opensearch.client.Client; import org.opensearch.common.unit.TimeValue; import org.opensearch.script.Script; +import org.opensearch.search.sort.FieldSortBuilder; +import org.opensearch.search.sort.SortOrder; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.legacy.domain.Field; import org.opensearch.sql.legacy.domain.KVValue; @@ -149,6 +144,12 @@ public void testIfScrollShouldBeOpenWithDifferentFormats() { queryAction.setFormat(Format.JDBC); queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setSize(settingFetchSize); + Mockito.verify(mockRequestBuilder).addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC); + Mockito.verify(mockRequestBuilder, never()).setScroll(timeValue); + + // Verify setScroll when SQL_PAGINATION_API_SEARCH_AFTER is set to false + mockLocalClusterStateAndIntializeMetricsForScroll(timeValue); + queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setScroll(timeValue); } @@ -168,6 +169,12 @@ public void testIfScrollShouldBeOpen() { mockLocalClusterStateAndInitializeMetrics(timeValue); queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setSize(settingFetchSize); + Mockito.verify(mockRequestBuilder).addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC); + Mockito.verify(mockRequestBuilder, never()).setScroll(timeValue); + + // Verify setScroll when SQL_PAGINATION_API_SEARCH_AFTER is set to false + mockLocalClusterStateAndIntializeMetricsForScroll(timeValue); + queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setScroll(timeValue); } @@ -195,6 +202,12 @@ public void testIfScrollShouldBeOpenWithDifferentFetchSize() { doReturn(mockRequestBuilder).when(mockRequestBuilder).setSize(userFetchSize); queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setSize(20); + Mockito.verify(mockRequestBuilder).addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC); + Mockito.verify(mockRequestBuilder, never()).setScroll(timeValue); + + // Verify setScroll when SQL_PAGINATION_API_SEARCH_AFTER is set to false + mockLocalClusterStateAndIntializeMetricsForScroll(timeValue); + queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setScroll(timeValue); } @@ -216,7 +229,9 @@ public void testIfScrollShouldBeOpenWithDifferentValidFetchSizeAndLimit() { queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setSize(userFetchSize); - Mockito.verify(mockRequestBuilder).setScroll(timeValue); + Mockito.verify(mockRequestBuilder).addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC); + // Skip setScroll when SQL_PAGINATION_API_SEARCH_AFTER is set to false + Mockito.verify(mockRequestBuilder, never()).setScroll(timeValue); /** fetchSize > LIMIT - no scroll */ userFetchSize = 5000; @@ -226,6 +241,14 @@ public void testIfScrollShouldBeOpenWithDifferentValidFetchSizeAndLimit() { queryAction.checkAndSetScroll(); Mockito.verify(mockRequestBuilder).setSize(limit); Mockito.verify(mockRequestBuilder, never()).setScroll(timeValue); + + // Verify setScroll when SQL_PAGINATION_API_SEARCH_AFTER is set to false + mockLocalClusterStateAndIntializeMetricsForScroll(timeValue); + /** fetchSize <= LIMIT - open scroll */ + userFetchSize = 1500; + doReturn(userFetchSize).when(mockSqlRequest).fetchSize(); + queryAction.checkAndSetScroll(); + Mockito.verify(mockRequestBuilder).setScroll(timeValue); } private void mockLocalClusterStateAndInitializeMetrics(TimeValue time) { @@ -236,6 +259,24 @@ private void mockLocalClusterStateAndInitializeMetrics(TimeValue time) { .when(mockLocalClusterState) .getSettingValue(Settings.Key.METRICS_ROLLING_WINDOW); doReturn(2L).when(mockLocalClusterState).getSettingValue(Settings.Key.METRICS_ROLLING_INTERVAL); + doReturn(true) + .when(mockLocalClusterState) + .getSettingValue(Settings.Key.SQL_PAGINATION_API_SEARCH_AFTER); + + Metrics.getInstance().registerDefaultMetrics(); + } + + private void mockLocalClusterStateAndIntializeMetricsForScroll(TimeValue time) { + LocalClusterState mockLocalClusterState = mock(LocalClusterState.class); + LocalClusterState.state(mockLocalClusterState); + doReturn(time).when(mockLocalClusterState).getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); + doReturn(3600L) + .when(mockLocalClusterState) + .getSettingValue(Settings.Key.METRICS_ROLLING_WINDOW); + doReturn(2L).when(mockLocalClusterState).getSettingValue(Settings.Key.METRICS_ROLLING_INTERVAL); + doReturn(false) + .when(mockLocalClusterState) + .getSettingValue(Settings.Key.SQL_PAGINATION_API_SEARCH_AFTER); Metrics.getInstance().registerDefaultMetrics(); }