Fix unit tests

Signed-off-by: Andre Kurait <[email protected]>
opensearch-project · Sep 4, 2024 · 4d23310 · 4d23310
1 parent d14aa7b
commit 4d23310
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/RFS/src/test/java/com/rfs/common/OpenSearchClientTest.java b/RFS/src/test/java/com/rfs/common/OpenSearchClientTest.java
@@ -137,7 +137,7 @@ void testBulkRequest_succeedAfterRetries() {
         var server500 = new HttpResponse(500, "", null, "{\"error\":\"Cannot Process Error!\"}");
 
         var restClient = mock(RestClient.class);
-        when(restClient.postAsync(any(), any(), any())).thenReturn(Mono.just(bothDocsFail))
+        when(restClient.postAsync(any(), any(), any(), any())).thenReturn(Mono.just(bothDocsFail))
             .thenReturn(Mono.just(oneFailure))
             .thenReturn(Mono.just(server500))
             .thenReturn(Mono.just(finalDocSuccess));
@@ -159,7 +159,7 @@ void testBulkRequest_succeedAfterRetries() {
         // Assertions
         // StepVerifier.create(responseMono).expectComplete().verify();
 
-        verify(restClient, times(4)).postAsync(any(), any(), any());
+        verify(restClient, times(4)).postAsync(any(), any(), any(), any());
         verifyNoInteractions(failedRequestLogger);
     }
 
@@ -169,7 +169,7 @@ void testBulkRequest_recordsTotalFailures() {
         var docFails = bulkItemResponse(true, List.of(itemEntryFailure(docId1)));
 
         var restClient = mock(RestClient.class);
-        when(restClient.postAsync(any(), any(), any())).thenReturn(Mono.just(docFails));
+        when(restClient.postAsync(any(), any(), any(), any())).thenReturn(Mono.just(docFails));
 
         var failedRequestLogger = mock(FailedRequestsLogger.class);
         var openSearchClient = spy(new OpenSearchClient(restClient, failedRequestLogger));
@@ -192,7 +192,7 @@ void testBulkRequest_recordsTotalFailures() {
         assertThat(exception.getMessage(), containsString("Retries exhausted"));
 
         var maxAttempts = maxRetries + 1;
-        verify(restClient, times(maxAttempts)).postAsync(any(), any(), any());
+        verify(restClient, times(maxAttempts)).postAsync(any(), any(), any(), any());
         verify(failedRequestLogger).logBulkFailure(any(), any(), any(), any());
         verifyNoMoreInteractions(failedRequestLogger);
     }

diff --git a/RFS/src/test/java/com/rfs/common/http/GzipPayloadRequestTransformerTest.java b/RFS/src/test/java/com/rfs/common/http/GzipPayloadRequestTransformerTest.java
@@ -39,7 +39,7 @@ void testGzipCompression(int size) throws Exception {
         // Generate test data
         ByteBuffer inputBuffer = generateTestData(size);
         Map<String, List<String>> headers = new HashMap<>();
-
+        headers.put("content-encoding", List.of("gzip"));
         // Store initial position and limit
         int initialPosition = inputBuffer.position();
         int initialLimit = inputBuffer.limit();
@@ -82,6 +82,7 @@ void testGzipCompressionWithDirectBuffer(int size) throws Exception {
         inputBuffer.flip();
 
         Map<String, List<String>> headers = new HashMap<>();
+        headers.put("content-encoding", List.of("gzip"));
 
         // Store initial position and limit
         int initialPosition = inputBuffer.position();
@@ -133,6 +134,7 @@ void testLargeInput() throws Exception {
         int largeSize = 50 * 1024 * 1024; // 50MB
         ByteBuffer largeBuffer = generateTestData(largeSize);
         Map<String, List<String>> headers = new HashMap<>();
+        headers.put("content-encoding", List.of("gzip"));
 
         // Store initial position and limit
         int initialPosition = largeBuffer.position();
@@ -167,6 +169,7 @@ public void testLargeInputWithDirectBuffer() throws Exception {
         largeBuffer.flip();
 
         Map<String, List<String>> headers = new HashMap<>();
+        headers.put("content-encoding", List.of("gzip"));
 
         // Store initial position and limit
         int initialPosition = largeBuffer.position();
@@ -208,8 +211,11 @@ private ByteBuffer generateTestData(int size) {
     }
 
     private byte[] decompress(ByteBuffer compressedBuffer) throws Exception {
+        ByteBuffer readArray = compressedBuffer.duplicate();
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
-        try (GZIPInputStream gzipInputStream = new GZIPInputStream(new ByteArrayInputStream(compressedBuffer.array()))) {
+        byte[] compressedArray = new byte[readArray.remaining()];
+        readArray.get(compressedArray);
+        try (GZIPInputStream gzipInputStream = new GZIPInputStream(new ByteArrayInputStream(compressedArray))) {
             byte[] buffer = new byte[1024];
             int len;
             while ((len = gzipInputStream.read(buffer)) > 0) {