From c23fb849130bb7d3c7ef7106829537442edbc664 Mon Sep 17 00:00:00 2001
From: Jai Balani <jbalani@jbalani-mn1.linkedin.biz>
Date: Mon, 16 Dec 2024 18:53:45 +0530
Subject: [PATCH] Added tests for state build for filecopy based replication

---
 .../ambry/store/StorageManagerTest.java       | 57 +++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/ambry-store/src/test/java/com/github/ambry/store/StorageManagerTest.java b/ambry-store/src/test/java/com/github/ambry/store/StorageManagerTest.java
index b2e86d054a..ab95636f00 100644
--- a/ambry-store/src/test/java/com/github/ambry/store/StorageManagerTest.java
+++ b/ambry-store/src/test/java/com/github/ambry/store/StorageManagerTest.java
@@ -68,6 +68,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -269,6 +270,62 @@ public void scheduleAndControlCompactionTest() throws Exception {
     shutdownAndAssertStoresInaccessible(storageManager, replicas);
   }
 
+  public MockId addRandomBlobToStore(Store store, long size, long expiresAtMs) throws StoreException {
+    final Random random = new Random();
+    short accountId = Utils.getRandomShort(TestUtils.RANDOM);
+    short containerId = Utils.getRandomShort(TestUtils.RANDOM);
+    short lifeVersion = MessageInfo.LIFE_VERSION_FROM_FRONTEND;
+
+    MockId id = new MockId(TestUtils.getRandomString(MOCK_ID_STRING_LENGTH), accountId, containerId);
+    long crc = random.nextLong();
+    MessageInfo info =
+        new MessageInfo(id, size, false, false, false, expiresAtMs, crc, id.getAccountId(), id.getContainerId(),
+            Utils.Infinite_Time, lifeVersion);
+    ByteBuffer buffer = ByteBuffer.wrap(TestUtils.getRandomBytes((int) size));
+    store.put(new MockMessageWriteSet(Collections.singletonList(info), Collections.singletonList(buffer)));
+    return id;
+  }
+
+  /**
+   * Test buildStateForFileCopy with newly created {@link ReplicaId}.
+   */
+  @Test
+  public void buildStateForFileCopyTest() throws Exception {
+    generateConfigs(true, false);
+    MockDataNodeId localNode = clusterMap.getDataNodes().get(0);
+    int newMountPathIndex = 3;
+    // add new MountPath to local node
+    File f = File.createTempFile("ambry", ".tmp");
+    File mountFile =
+        new File(f.getParent(), "mountpathfile" + MockClusterMap.PLAIN_TEXT_PORT_START_NUMBER + newMountPathIndex);
+    MockClusterMap.deleteFileOrDirectory(mountFile);
+    assertTrue("Couldn't create mount path directory", mountFile.mkdir());
+    localNode.addMountPaths(Collections.singletonList(mountFile.getAbsolutePath()));
+    int newPartitionId = 803;
+    PartitionId newPartition =
+        new MockPartitionId(newPartitionId, MockClusterMap.DEFAULT_PARTITION_CLASS, clusterMap.getDataNodes(), newMountPathIndex);
+    StorageManager storageManager = createStorageManager(localNode, metricRegistry, null);
+    storageManager.start();
+    // test add store onto a new disk, which should succeed
+    assertTrue("Add new store should succeed", storageManager.addBlobStore(newPartition.getReplicaIds().get(0)));
+    assertNotNull("The store shouldn't be null because new store is successfully added",
+        storageManager.getStore(newPartition, false));
+    DiskManager dm = storageManager.getDiskManager(newPartition);
+    Store store = dm.getStore(newPartition, false);
+    MockId id1 = addRandomBlobToStore(store, 100, Utils.Infinite_Time);
+    MockId id2 = addRandomBlobToStore(store, 200, Utils.Infinite_Time);
+    // Create storage manager again to create disk managers again
+    storageManager.shutdown();
+    storageManager = createStorageManager(localNode, metricRegistry, null);
+    storageManager.start();
+    storageManager.buildStateForFileCopy(newPartition.getReplicaIds().get(0));
+    dm = storageManager.getDiskManager(newPartition);
+    store = dm.getStore(newPartition, false);
+    assertNotNull(store.get(Collections.singletonList(id1), EnumSet.noneOf(StoreGetOptions.class)));
+    assertNotNull(store.get(Collections.singletonList(id2), EnumSet.noneOf(StoreGetOptions.class)));
+  }
+
+
   /**
    * Test add new BlobStore with given {@link ReplicaId}.
    */