Fix failing tests

Signed-off-by: Bhumika Saini <sabhumik@amazon.com>
opensearch-project · Jun 8, 2023 · d34ef74 · d34ef74
1 parent 9850c65
commit d34ef74
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 44 deletions.
diff --git a/server/src/main/java/org/opensearch/index/translog/transfer/TranslogTransferManager.java b/server/src/main/java/org/opensearch/index/translog/transfer/TranslogTransferManager.java
@@ -206,27 +206,39 @@ private TransferFileSnapshot prepareMetadata(TransferSnapshot transferSnapshot)
             );
         TranslogTransferMetadata translogTransferMetadata = transferSnapshot.getTranslogTransferMetadata();
         translogTransferMetadata.setGenerationToPrimaryTermMapper(new HashMap<>(generationPrimaryTermMap));
+
+        return new TransferFileSnapshot(
+            getFileName(translogTransferMetadata.getPrimaryTerm(), translogTransferMetadata.getGeneration()),
+            getMetadataBytes(translogTransferMetadata),
+            translogTransferMetadata.getPrimaryTerm()
+        );
+    }
+
+    /**
+     * Get the metadata bytes for a {@link TranslogTransferMetadata} object
+     *
+     * @param metadata The object to be parsed
+     * @return Byte representation for the given metadata
+     * @throws IOException
+     */
+    public byte[] getMetadataBytes(TranslogTransferMetadata metadata) throws IOException {
         byte[] metadataBytes;
 
         try (BytesStreamOutput output = new BytesStreamOutput()) {
             try (
                 OutputStreamIndexOutput indexOutput = new OutputStreamIndexOutput(
-                    "translog transfer metadata " + translogTransferMetadata.getPrimaryTerm(),
-                    getFileName(translogTransferMetadata.getPrimaryTerm(), translogTransferMetadata.getGeneration()),
+                    "translog transfer metadata " + metadata.getPrimaryTerm(),
+                    getFileName(metadata.getPrimaryTerm(), metadata.getGeneration()),
                     output,
                     TranslogTransferMetadata.BUFFER_SIZE
                 )
             ) {
-                metadataStreamWrapper.writeStream(indexOutput, translogTransferMetadata);
+                metadataStreamWrapper.writeStream(indexOutput, metadata);
             }
             metadataBytes = BytesReference.toBytes(output.bytes());
         }
 
-        return new TransferFileSnapshot(
-            getFileName(translogTransferMetadata.getPrimaryTerm(), translogTransferMetadata.getGeneration()),
-            metadataBytes,
-            translogTransferMetadata.getPrimaryTerm()
-        );
+        return metadataBytes;
     }
 
     /**

diff --git a/...src/main/java/org/opensearch/index/translog/transfer/TranslogTransferMetadataHandler.java b/...src/main/java/org/opensearch/index/translog/transfer/TranslogTransferMetadataHandler.java
@@ -16,6 +16,11 @@
 import java.util.HashMap;
 import java.util.Map;
 
+/**
+ * Handler for {@link TranslogTransferMetadata}
+ *
+ * @opensearch.internal
+ */
 public class TranslogTransferMetadataHandler implements IndexIOStreamHandler<TranslogTransferMetadata> {
 
     /**
@@ -54,7 +59,5 @@ public void writeContent(IndexOutput indexOutput, TranslogTransferMetadata conte
         } else {
             indexOutput.writeMapOfStrings(new HashMap<>());
         }
-
-        indexOutput.close();
     }
 }
diff --git a/...er/src/test/java/org/opensearch/index/translog/transfer/TranslogTransferManagerTests.java b/...er/src/test/java/org/opensearch/index/translog/transfer/TranslogTransferManagerTests.java
@@ -8,15 +8,12 @@
 
 package org.opensearch.index.translog.transfer;
 
-import org.apache.lucene.store.OutputStreamIndexOutput;
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.mockito.Mockito;
 import org.opensearch.action.ActionListener;
 import org.opensearch.common.blobstore.BlobContainer;
 import org.opensearch.common.blobstore.BlobPath;
 import org.opensearch.common.blobstore.BlobStore;
-import org.opensearch.common.bytes.BytesReference;
-import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.util.set.Sets;
 import org.opensearch.index.shard.ShardId;
 import org.opensearch.index.translog.Translog;
@@ -34,7 +31,6 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -46,7 +42,6 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
-import static org.opensearch.index.translog.transfer.TranslogTransferMetadata.getFileName;
 
 @LuceneTestCase.SuppressFileSystems("*")
 public class TranslogTransferManagerTests extends OpenSearchTestCase {
@@ -209,7 +204,7 @@ public void testReadMetadataSingleFile() throws IOException {
 
         TranslogTransferMetadata metadata = createTransferSnapshot().getTranslogTransferMetadata();
         when(transferService.downloadBlob(any(BlobPath.class), eq("12__234"))).thenReturn(
-            new ByteArrayInputStream(getMetadataBytes(metadata))
+            new ByteArrayInputStream(translogTransferManager.getMetadataBytes(metadata))
         );
 
         assertEquals(metadata, translogTransferManager.readMetadata());
@@ -227,7 +222,7 @@ public void testReadMetadataMultipleFiles() throws IOException {
 
         TranslogTransferMetadata metadata = createTransferSnapshot().getTranslogTransferMetadata();
         when(transferService.downloadBlob(any(BlobPath.class), eq("12__235"))).thenReturn(
-            new ByteArrayInputStream(getMetadataBytes(metadata))
+            new ByteArrayInputStream(translogTransferManager.getMetadataBytes(metadata))
         );
 
         assertEquals(metadata, translogTransferManager.readMetadata());
@@ -386,31 +381,4 @@ public void testDeleteTranslogFailure() throws Exception {
         translogTransferManager.deleteGenerationAsync(primaryTerm, Set.of(19L), () -> {});
         assertEquals(2, tracker.allUploaded().size());
     }
-
-    private byte[] getMetadataBytes(TranslogTransferMetadata metadata) throws IOException {
-        byte[] metadataBytes;
-
-        try (BytesStreamOutput output = new BytesStreamOutput()) {
-            try (
-                OutputStreamIndexOutput indexOutput = new OutputStreamIndexOutput(
-                    "translog transfer metadata " + metadata.getPrimaryTerm(),
-                    getFileName(metadata.getPrimaryTerm(), metadata.getGeneration()),
-                    output,
-                    TranslogTransferMetadata.BUFFER_SIZE
-                )
-            ) {
-                indexOutput.writeLong(metadata.getPrimaryTerm());
-                indexOutput.writeLong(metadata.getGeneration());
-                indexOutput.writeLong(metadata.getMinTranslogGeneration());
-                if (metadata.getGenerationToPrimaryTermMapper() != null) {
-                    indexOutput.writeMapOfStrings(metadata.getGenerationToPrimaryTermMapper());
-                } else {
-                    indexOutput.writeMapOfStrings(new HashMap<>());
-                }
-            }
-            metadataBytes = BytesReference.toBytes(output.bytes());
-        }
-
-        return metadataBytes;
-    }
 }