Skip to content

Commit

Permalink
FIx unit tests for PIT changes
Browse files Browse the repository at this point in the history 
Signed-off-by: Manasvini B S <[email protected]>
  • Loading branch information
@manasvinibs
manasvinibs committed Aug 15, 2024
1 parent 8c8aba8 commit 167dce0
Show file tree
Hide file tree
Showing 4 changed files with 204 additions and 35 deletions.
3 changes: 2 additions & 1 deletion legacy/src/main/java/org/opensearch/sql/legacy/pit/PointInTimeHandlerImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ public boolean create() {
client.createPit(
createPitRequest,
new ActionListener<>() {

@Override
public void onResponse(CreatePitResponse createPitResponse) {
pitId = createPitResponse.getId();
Expand Down Expand Up @@ -112,14 +113,14 @@ public void onFailure(Exception e) {
LOG.error("Error occurred while deleting PIT", e);
}
});

while (deleteStatus == null) {
try {
Thread.sleep(100);
} catch (InterruptedException e) {
LOG.error("Error occurred while deleting PIT", e);
}
}

return deleteStatus;
}
}
115 changes: 93 additions & 22 deletions legacy/src/test/java/org/opensearch/sql/legacy/pit/PointInTimeHandlerImplTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,64 +2,135 @@
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.sql.legacy.pit;

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

import java.util.Collections;
import java.util.concurrent.CompletableFuture;
import org.junit.Before;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.opensearch.action.search.CreatePitResponse;
import org.opensearch.action.search.DeletePitResponse;
import org.opensearch.client.Client;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.sql.common.setting.Settings;
import org.opensearch.sql.legacy.esdomain.LocalClusterState;
import org.opensearch.sql.opensearch.setting.OpenSearchSettings;

public class PointInTimeHandlerImplTest {

@Mock private Client mockClient;
private String[] indices = {"index1", "index2"};
private PointInTimeHandlerImpl pointInTimeHandlerImpl;
@Captor private ArgumentCaptor<ActionListener<DeletePitResponse>> listenerCaptor;
@Captor private ArgumentCaptor<ActionListener<DeletePitResponse>> listenerCaptorForDelete;
@Captor private ArgumentCaptor<ActionListener<CreatePitResponse>> listenerCaptorForCreate;
private final String PIT_ID = "testId";
private CreatePitResponse mockCreatePitResponse;
private CompletableFuture<CreatePitResponse> completableFuture;
private CompletableFuture<DeletePitResponse> completableFutureForDelete;
private Exception exception;
private DeletePitResponse mockDeletePitResponse;

@Mock private OpenSearchSettings settings;

/*@Before
@Before
public void setUp() {
MockitoAnnotations.initMocks(this);
pointInTimeHandlerImpl = new PointInTimeHandlerImpl(mockClient, indices);
}

@Test
public void testCreate() {
when(LocalClusterState.state().getSettingValue(SQL_CURSOR_KEEP_ALIVE))
doReturn(Collections.emptyList()).when(settings).getSettings();
when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE))
.thenReturn(new TimeValue(10000));
LocalClusterState.state().setPluginSettings(settings);

CreatePitResponse mockCreatePitResponse = mock(CreatePitResponse.class);
mockCreatePitResponse = mock(CreatePitResponse.class);
mockDeletePitResponse = mock(DeletePitResponse.class);
RestStatus mockRestStatus = mock(RestStatus.class);
when(mockDeletePitResponse.status()).thenReturn(mockRestStatus);
when(mockDeletePitResponse.status().getStatus()).thenReturn(200);
when(mockCreatePitResponse.getId()).thenReturn(PIT_ID);

CompletableFuture<CreatePitResponse> completableFuture =
CompletableFuture.completedFuture(mockCreatePitResponse);
completableFuture = CompletableFuture.completedFuture(mockCreatePitResponse);
completableFutureForDelete = CompletableFuture.completedFuture(mockDeletePitResponse);
exception = mock(Exception.class);
}

@Test
public void testCreate() {
doAnswer(
invocation -> {
ActionListener<CreatePitResponse> actionListener = invocation.getArgument(1);
actionListener.onResponse(mockCreatePitResponse);
return completableFuture;
})
.when(mockClient)
.createPit(any(), any());
.createPit(any(), listenerCaptorForCreate.capture());

boolean status = pointInTimeHandlerImpl.create();
verify(mockClient).createPit(any(), listenerCaptorForCreate.capture());
listenerCaptorForCreate.getValue().onResponse(mockCreatePitResponse);
verify(mockCreatePitResponse, times(2)).getId();
assertTrue(status);
}

pointInTimeHandlerImpl.create();
@Test
public void testCreateForFailure() {
doAnswer(
invocation -> {
ActionListener<CreatePitResponse> actionListener = invocation.getArgument(1);
actionListener.onFailure(exception);
return completableFuture;
})
.when(mockClient)
.createPit(any(), listenerCaptorForCreate.capture());

assertEquals(PIT_ID, pointInTimeHandlerImpl.getPitId());
boolean status = pointInTimeHandlerImpl.create();
verify(mockClient).createPit(any(), listenerCaptorForCreate.capture());
listenerCaptorForCreate.getValue().onResponse(mockCreatePitResponse);
assertFalse(status);
}

@Test
public void testDelete() {
DeletePitResponse mockedResponse = mock(DeletePitResponse.class);
RestStatus mockRestStatus = mock(RestStatus.class);
when(mockedResponse.status()).thenReturn(mockRestStatus);
when(mockedResponse.status().getStatus()).thenReturn(200);
pointInTimeHandlerImpl.setPitId(PIT_ID);
pointInTimeHandlerImpl.delete();
verify(mockClient).deletePits(any(), listenerCaptor.capture());
listenerCaptor.getValue().onResponse(mockedResponse);
}*/
doAnswer(
invocation -> {
ActionListener<DeletePitResponse> actionListener = invocation.getArgument(1);
actionListener.onResponse(mockDeletePitResponse);
return completableFutureForDelete;
})
.when(mockClient)
.deletePits(any(), listenerCaptorForDelete.capture());

boolean status = pointInTimeHandlerImpl.delete();
assertTrue(status);
verify(mockClient).deletePits(any(), listenerCaptorForDelete.capture());
listenerCaptorForDelete.getValue().onResponse(mockDeletePitResponse);
}

@Test
public void testDeleteForFailure() {
doAnswer(
invocation -> {
ActionListener<DeletePitResponse> actionListener = invocation.getArgument(1);
actionListener.onFailure(exception);
return completableFutureForDelete;
})
.when(mockClient)
.deletePits(any(), listenerCaptorForDelete.capture());

boolean status = pointInTimeHandlerImpl.delete();
assertFalse(status);
verify(mockClient).deletePits(any(), listenerCaptorForDelete.capture());
listenerCaptorForDelete.getValue().onResponse(mockDeletePitResponse);
}
}
58 changes: 57 additions & 1 deletion legacy/src/test/java/org/opensearch/sql/legacy/unittest/cursor/DefaultCursorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand All @@ -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() + ":"));
}

Expand Down
63 changes: 52 additions & 11 deletions legacy/src/test/java/org/opensearch/sql/legacy/unittest/query/DefaultQueryActionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}

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

Expand Down Expand Up @@ -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);
}

Expand All @@ -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;
Expand All @@ -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) {
Expand All @@ -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();
}
Expand Down

0 comments on commit 167dce0

Please sign in to comment.