Skip to content

Commit

Permalink
Add UTs
Browse files Browse the repository at this point in the history
Signed-off-by: Arpit Bandejiya <abandeji@amazon.com>
  • Loading branch information
Arpit-Bandejiya committed Aug 27, 2024
1 parent 5994e9c commit f6d1b5a
Show file tree
Hide file tree
Showing 4 changed files with 139 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,7 @@ private PublishWithJoinResponse handleIncomingPublishRequest(BytesTransportReque
}

// package private for testing
PublishWithJoinResponse handleIncomingRemotePublishRequest(RemotePublishRequest request) throws IOException, RuntimeException {
PublishWithJoinResponse handleIncomingRemotePublishRequest(RemotePublishRequest request) throws IOException, IllegalStateException {
try {
if (transportService.getLocalNode().equals(request.getSourceNode())) {
return acceptRemoteStateOnLocalNode(request);
Expand Down Expand Up @@ -302,8 +302,7 @@ PublishWithJoinResponse handleIncomingRemotePublishRequest(RemotePublishRequest
}
} catch (Exception e) {
remoteClusterStateService.readMetadataFailed();
if (e instanceof IOException) throw new IOException("IOException in reading remote cluster state", e);
throw new RuntimeException("Runtime exception in reading remote cluster state", e);
throw e;
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,4 +33,11 @@ public void fullDownloadState() {
fullDownloadCount.incrementAndGet();
}

public long getDiffDownloadCount() {
return diffDownloadCount.get();
}

public long getFullDownloadCount() {
return fullDownloadCount.get();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,9 @@
import org.opensearch.common.settings.Settings;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.gateway.remote.ClusterMetadataManifest;
import org.opensearch.gateway.remote.ClusterStateDiffManifest;
import org.opensearch.gateway.remote.RemoteClusterStateService;
import org.opensearch.gateway.remote.RemoteDownloadStats;
import org.opensearch.node.Node;
import org.opensearch.telemetry.tracing.noop.NoopTracer;
import org.opensearch.test.OpenSearchTestCase;
Expand All @@ -62,8 +64,12 @@
import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.is;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

public class PublicationTransportHandlerTests extends OpenSearchTestCase {
Expand Down Expand Up @@ -160,7 +166,8 @@ public void testHandleIncomingRemotePublishRequestWhenNoCurrentPublishRequest()
() -> handler.handleIncomingRemotePublishRequest(remotePublishRequest)
);
assertThat(e.getMessage(), containsString("publication to self failed"));
Mockito.verifyNoInteractions(remoteClusterStateService);
verify(remoteClusterStateService, times(1)).readMetadataFailed();
verifyNoMoreInteractions(remoteClusterStateService);
}

public void testHandleIncomingRemotePublishRequestWhenTermMismatch() {
Expand All @@ -185,7 +192,8 @@ public void testHandleIncomingRemotePublishRequestWhenTermMismatch() {
() -> handler.handleIncomingRemotePublishRequest(remotePublishRequest)
);
assertThat(e.getMessage(), containsString("publication to self failed"));
Mockito.verifyNoInteractions(remoteClusterStateService);
verify(remoteClusterStateService, times(1)).readMetadataFailed();
verifyNoMoreInteractions(remoteClusterStateService);
}

public void testHandleIncomingRemotePublishRequestWhenVersionMismatch() {
Expand All @@ -210,7 +218,8 @@ public void testHandleIncomingRemotePublishRequestWhenVersionMismatch() {
() -> handler.handleIncomingRemotePublishRequest(remotePublishRequest)
);
assertThat(e.getMessage(), containsString("publication to self failed"));
Mockito.verifyNoInteractions(remoteClusterStateService);
verify(remoteClusterStateService, times(1)).readMetadataFailed();
verifyNoMoreInteractions(remoteClusterStateService);
}

public void testHandleIncomingRemotePublishRequestForLocalNode() throws IOException {
Expand All @@ -235,6 +244,119 @@ public void testHandleIncomingRemotePublishRequestForLocalNode() throws IOExcept
Mockito.verifyNoInteractions(remoteClusterStateService);
}

public void testDownloadRemotePersistedFailedStats() throws IOException {
RemoteDownloadStats remoteDownloadStats = new RemoteDownloadStats();
RemoteClusterStateService remoteClusterStateService = mock(RemoteClusterStateService.class);
when(remoteClusterStateService.getDownloadStats()).thenReturn(remoteDownloadStats);

doAnswer((i) -> {
remoteDownloadStats.stateFailed();
return null;
}).when(remoteClusterStateService).readMetadataFailed();

PublishWithJoinResponse expectedPublishResponse = new PublishWithJoinResponse(new PublishResponse(TERM, VERSION), Optional.empty());
Function<PublishRequest, PublishWithJoinResponse> handlePublishRequest = p -> expectedPublishResponse;
final PublicationTransportHandler handler = getPublicationTransportHandler(handlePublishRequest, remoteClusterStateService);
RemotePublishRequest remotePublishRequest = new RemotePublishRequest(
secondNode,
TERM,
VERSION,
CLUSTER_NAME,
CLUSTER_UUID,
MANIFEST_FILE
);
ClusterState clusterState = buildClusterState(TERM, VERSION);
PublishRequest publishRequest = new PublishRequest(clusterState);
handler.setCurrentPublishRequestToSelf(publishRequest);

assertThrows(IllegalStateException.class, () -> handler.handleIncomingRemotePublishRequest(remotePublishRequest));
assertEquals(1, remoteClusterStateService.getDownloadStats().getFailedCount());
}

public void testDownloadRemotePersistedDiffStats() throws IOException {
RemoteDownloadStats remoteDownloadStats = new RemoteDownloadStats();
RemoteClusterStateService remoteClusterStateService = mock(RemoteClusterStateService.class);
when(remoteClusterStateService.getDownloadStats()).thenReturn(remoteDownloadStats);
ClusterMetadataManifest metadataManifest = new ClusterMetadataManifest.Builder().diffManifest(
new ClusterStateDiffManifest.Builder().fromStateUUID("state-uuid").build()
).build();
when(remoteClusterStateService.getClusterMetadataManifestByFileName(any(), any())).thenReturn(metadataManifest);

doAnswer((i) -> {
remoteDownloadStats.diffDownloadState();
return null;
}).when(remoteClusterStateService).diffDownloadState();

doAnswer((i) -> {
remoteDownloadStats.fullDownloadState();
return null;
}).when(remoteClusterStateService).fullDownloadState();

PublishWithJoinResponse expectedPublishResponse = new PublishWithJoinResponse(new PublishResponse(TERM, VERSION), Optional.empty());
Function<PublishRequest, PublishWithJoinResponse> handlePublishRequest = p -> expectedPublishResponse;
final PublicationTransportHandler handler = getPublicationTransportHandler(handlePublishRequest, remoteClusterStateService);
ClusterState clusterState = mock(ClusterState.class);
handler.setLastSeenClusterState(clusterState);
when(clusterState.stateUUID()).thenReturn("state-uuid");

RemotePublishRequest remotePublishRequest = new RemotePublishRequest(
secondNode,
TERM,
VERSION,
CLUSTER_NAME,
CLUSTER_UUID,
MANIFEST_FILE
);
clusterState = buildClusterState(TERM, VERSION);
PublishRequest publishRequest = new PublishRequest(clusterState);
handler.setCurrentPublishRequestToSelf(publishRequest);
assertThrows(NullPointerException.class, () -> handler.handleIncomingRemotePublishRequest(remotePublishRequest));
assertEquals(1, remoteDownloadStats.getDiffDownloadCount());
assertEquals(0, remoteDownloadStats.getFullDownloadCount());
}

public void testDownloadRemotePersistedFullStats() throws IOException {
RemoteDownloadStats remoteDownloadStats = new RemoteDownloadStats();
RemoteClusterStateService remoteClusterStateService = mock(RemoteClusterStateService.class);
when(remoteClusterStateService.getDownloadStats()).thenReturn(remoteDownloadStats);
ClusterMetadataManifest metadataManifest = new ClusterMetadataManifest.Builder().diffManifest(
new ClusterStateDiffManifest.Builder().fromStateUUID("state-uuid2").build()
).build();
when(remoteClusterStateService.getClusterMetadataManifestByFileName(any(), any())).thenReturn(metadataManifest);

doAnswer((i) -> {
remoteDownloadStats.diffDownloadState();
return null;
}).when(remoteClusterStateService).diffDownloadState();

doAnswer((i) -> {
remoteDownloadStats.fullDownloadState();
return null;
}).when(remoteClusterStateService).fullDownloadState();

PublishWithJoinResponse expectedPublishResponse = new PublishWithJoinResponse(new PublishResponse(TERM, VERSION), Optional.empty());
Function<PublishRequest, PublishWithJoinResponse> handlePublishRequest = p -> expectedPublishResponse;
final PublicationTransportHandler handler = getPublicationTransportHandler(handlePublishRequest, remoteClusterStateService);
ClusterState clusterState = mock(ClusterState.class);
handler.setLastSeenClusterState(clusterState);
when(clusterState.stateUUID()).thenReturn("state-uuid");

RemotePublishRequest remotePublishRequest = new RemotePublishRequest(
secondNode,
TERM,
VERSION,
CLUSTER_NAME,
CLUSTER_UUID,
MANIFEST_FILE
);
clusterState = buildClusterState(TERM, VERSION);
PublishRequest publishRequest = new PublishRequest(clusterState);
handler.setCurrentPublishRequestToSelf(publishRequest);
assertThrows(NullPointerException.class, () -> handler.handleIncomingRemotePublishRequest(remotePublishRequest));
assertEquals(0, remoteDownloadStats.getDiffDownloadCount());
assertEquals(1, remoteDownloadStats.getFullDownloadCount());
}

public void testHandleIncomingRemotePublishRequestWhenManifestNotFound() throws IOException {
RemoteClusterStateService remoteClusterStateService = mock(RemoteClusterStateService.class);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -851,6 +851,10 @@ public void testGetClusterStateForManifest_IncludeEphemeral() throws IOException
when(mockedResult.getComponent()).thenReturn(COORDINATION_METADATA);
RemoteClusterStateService mockService = spy(remoteClusterStateService);
mockService.getClusterStateForManifest(ClusterName.DEFAULT.value(), manifest, NODE_ID, true);

assertNotNull(remoteClusterStateService.getDownloadStats());
assertEquals(1, remoteClusterStateService.getDownloadStats().getSuccessCount());
assertEquals(0, remoteClusterStateService.getDownloadStats().getFailedCount());
verify(mockService, times(1)).readClusterStateInParallel(
any(),
eq(manifest),
Expand Down Expand Up @@ -2568,7 +2572,7 @@ public void testGetValidPreviousClusterUUIDWhenLastUUIDUncommitted() throws IOEx
assertThat(previousClusterUUID, equalTo("cluster-uuid2"));
}

public void testRemoteStateStats() throws IOException {
public void testRemoteStateUploadStats() throws IOException {
final ClusterState clusterState = generateClusterStateWithOneIndex().nodes(nodesWithLocalNodeClusterManager()).build();
mockBlobStoreObjects();
remoteClusterStateService.start();
Expand Down

0 comments on commit f6d1b5a

Please sign in to comment.