FIx unit tests for PIT changes

Signed-off-by: Manasvini B S <[email protected]>
manasvinibs · Aug 14, 2024 · fe45ecd · fe45ecd
1 parent d022def
commit fe45ecd
Show file tree

Hide file tree

Showing 2 changed files with 109 additions and 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
@@ -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,14 +58,37 @@ 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+//");
     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() + ":"));
   }
 

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
@@ -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();
   }