Add unit tests for Global RT DIV state methods to fix diff coverage

Add tests covering the new branches introduced for Global RT DIV state
storage and retrieval:

- RocksDBStorageEngineTest: tests for putGlobalRtDivState,
  getGlobalRtDivState, and clearGlobalRtDivState happy paths and
  illegal argument branches (negative partition id, metadata partition id)
- LeaderFollowerStoreIngestionTaskTest: tests for
  putGlobalRtDivStateInMetadata (invalid key prefix exception and valid
  delegation) and readGlobalRtDivState (metadata-present early return,
  metadata-absent fallthrough, and non-prefix key skip)

Diff branch coverage for da-vinci-client goes from 11% to 67.65%,
exceeding the 45% minimum threshold.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: KaiSernLim

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTaskTest.java

@@ -37,6 +37,7 @@
 import com.linkedin.davinci.stats.ingestion.heartbeat.HeartbeatMonitoringService;
 import com.linkedin.davinci.storage.StorageMetadataService;
 import com.linkedin.davinci.storage.StorageService;
+import com.linkedin.davinci.storage.chunking.ChunkedValueManifestContainer;
 import com.linkedin.davinci.store.DelegatingStorageEngine;
 import com.linkedin.davinci.store.view.MaterializedViewWriter;
 import com.linkedin.davinci.store.view.VeniceViewWriter;
@@ -45,6 +46,7 @@
 import com.linkedin.davinci.validation.PartitionTracker;
 import com.linkedin.venice.compression.CompressionStrategy;
 import com.linkedin.venice.compression.VeniceCompressor;
+import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.exceptions.VeniceTimeoutException;
 import com.linkedin.venice.kafka.protocol.ControlMessage;
 import com.linkedin.venice.kafka.protocol.Delete;
@@ -96,6 +98,7 @@
 import com.linkedin.venice.writer.VeniceWriter;
 import it.unimi.dsi.fastutil.objects.Object2IntMaps;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.HashMap;
@@ -1397,4 +1400,97 @@ public void testDolStampIsDolComplete() {
     dolStampIncomplete.setDolConsumed(true);
     assertFalse(dolStampIncomplete.isDolComplete());
   }
+
+  /**
+   * Tests that {@link StoreIngestionTask#putGlobalRtDivStateInMetadata} throws a VeniceException
+   * when the key does not start with the expected prefix, and delegates to storageMetadataService
+   * when the key is valid.
+   */
+  @Test
+  public void testPutGlobalRtDivStateInMetadata() throws Exception {
+    LeaderFollowerStoreIngestionTask ingestionTask = mock(LeaderFollowerStoreIngestionTask.class);
+    doCallRealMethod().when(ingestionTask).putGlobalRtDivStateInMetadata(anyInt(), any(), any());
+
+    Put put = new Put();
+    put.putValue = ByteBuffer.wrap("test-value".getBytes());
+
+    // Invalid key (missing prefix) should throw VeniceException immediately, before any field access
+    byte[] invalidKey = "INVALID_KEY.localhost:9092".getBytes();
+    Assert.assertThrows(VeniceException.class, () -> ingestionTask.putGlobalRtDivStateInMetadata(0, invalidKey, put));
+
+    // Valid key should call storageMetadataService.putGlobalRtDivState
+    StorageMetadataService mockSms = mock(StorageMetadataService.class);
+    String versionTopic = "testStore_v1";
+    injectField(ingestionTask, StoreIngestionTask.class, "storageMetadataService", mockSms);
+    injectField(ingestionTask, StoreIngestionTask.class, "kafkaVersionTopic", versionTopic);
+
+    String brokerUrl = "localhost:9092";
+    byte[] validKey = (StoreIngestionTask.GLOBAL_RT_DIV_KEY_PREFIX + brokerUrl).getBytes();
+    ingestionTask.putGlobalRtDivStateInMetadata(0, validKey, put);
+
+    ArgumentCaptor<byte[]> valueCaptor = ArgumentCaptor.forClass(byte[].class);
+    verify(mockSms, times(1)).putGlobalRtDivState(eq(versionTopic), eq(0), eq(brokerUrl), valueCaptor.capture());
+  }
+
+  /**
+   * Tests the new metadata-first lookup in {@link LeaderFollowerStoreIngestionTask#readGlobalRtDivState}:
+   * - When metadata storage has a value, it is returned without hitting legacy storage.
+   * - When metadata storage is empty, the method falls through to the legacy path.
+   * - When the key does not start with the GLOBAL_RT_DIV_KEY_PREFIX, the metadata lookup is skipped.
+   */
+  @Test
+  public void testReadGlobalRtDivStateMetadataPath() throws Exception {
+    LeaderFollowerStoreIngestionTask ingestionTask = mock(LeaderFollowerStoreIngestionTask.class);
+    doCallRealMethod().when(ingestionTask)
+        .readGlobalRtDivState(any(), anyInt(), any(), any(ChunkedValueManifestContainer.class));
+
+    StorageMetadataService mockSms = mock(StorageMetadataService.class);
+    String versionTopic = "testStore_v1";
+    InternalAvroSpecificSerializer<GlobalRtDivState> serializer =
+        AvroProtocolDefinition.GLOBAL_RT_DIV_STATE.getSerializer();
+
+    injectField(ingestionTask, StoreIngestionTask.class, "storageMetadataService", mockSms);
+    injectField(ingestionTask, StoreIngestionTask.class, "kafkaVersionTopic", versionTopic);
+    injectField(ingestionTask, LeaderFollowerStoreIngestionTask.class, "globalRtDivStateSerializer", serializer);
+
+    String brokerUrl = "localhost:9092";
+    byte[] keyBytes = (StoreIngestionTask.GLOBAL_RT_DIV_KEY_PREFIX + brokerUrl).getBytes();
+    PubSubTopicPartition topicPartition = mock(PubSubTopicPartition.class);
+    doReturn(0).when(topicPartition).getPartitionNumber();
+    ChunkedValueManifestContainer manifestContainer = new ChunkedValueManifestContainer();
+
+    // Serialize a real GlobalRtDivState for the round-trip test
+    GlobalRtDivState expectedState =
+        new GlobalRtDivState(brokerUrl, Collections.emptyMap(), InMemoryPubSubPosition.of(5).toWireFormatBuffer());
+    byte[] serializedState = serializer.serialize(null, expectedState);
+
+    // Case 1: metadata present → returns deserialized state without touching legacy storage
+    doReturn(Optional.of(serializedState)).when(mockSms).getGlobalRtDivState(versionTopic, 0, brokerUrl);
+    GlobalRtDivState result =
+        ingestionTask.readGlobalRtDivState(keyBytes, GLOBAL_RT_DIV_VERSION, topicPartition, manifestContainer);
+    assertNotNull(result);
+    assertEquals(result.srcUrl.toString(), brokerUrl);
+
+    // Case 2: metadata absent → falls through to legacy; legacy returns null (compressor field is null on mock),
+    // so the method returns null
+    doReturn(Optional.empty()).when(mockSms).getGlobalRtDivState(versionTopic, 0, brokerUrl);
+    GlobalRtDivState fallThroughResult =
+        ingestionTask.readGlobalRtDivState(keyBytes, GLOBAL_RT_DIV_VERSION, topicPartition, manifestContainer);
+    Assert.assertNull(fallThroughResult);
+
+    // Case 3: key does not start with the prefix → metadata lookup skipped, goes straight to legacy (returns null)
+    byte[] nonPrefixKey = "REGULAR_KEY.localhost:9092".getBytes();
+    GlobalRtDivState nonPrefixResult =
+        ingestionTask.readGlobalRtDivState(nonPrefixKey, GLOBAL_RT_DIV_VERSION, topicPartition, manifestContainer);
+    Assert.assertNull(nonPrefixResult);
+    // storageMetadataService.getGlobalRtDivState should NOT have been called for the non-prefix key
+    verify(mockSms, times(0)).getGlobalRtDivState(versionTopic, 0, "REGULAR_KEY.localhost:9092".substring(0));
+  }
+
+  private static void injectField(Object target, Class<?> declaringClass, String fieldName, Object value)
+      throws Exception {
+    Field field = declaringClass.getDeclaredField(fieldName);
+    field.setAccessible(true);
+    field.set(target, value);
+  }
 }

clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStorageEngineTest.java

@@ -288,4 +288,59 @@ public void testHasConflictPersistedStoreEngineConfig() {
     Assert.assertTrue(result);
 
   }
+
+  @Test
+  public void testGetAndPutGlobalRtDivState() {
+    RocksDBStorageEngine rocksDBStorageEngine = (RocksDBStorageEngine) getTestStoreEngine();
+    int partitionId = PARTITION_ID;
+    String brokerUrl = "localhost:9092";
+    byte[] valueBytes = "test-global-rt-div-value".getBytes();
+
+    // Put and get round-trip
+    rocksDBStorageEngine.putGlobalRtDivState(partitionId, brokerUrl, valueBytes);
+    Assert.assertTrue(rocksDBStorageEngine.getGlobalRtDivState(partitionId, brokerUrl).isPresent());
+    Assert.assertEquals(rocksDBStorageEngine.getGlobalRtDivState(partitionId, brokerUrl).get(), valueBytes);
+
+    // Clear and verify absent
+    rocksDBStorageEngine.clearGlobalRtDivState(partitionId, brokerUrl);
+    Assert.assertFalse(rocksDBStorageEngine.getGlobalRtDivState(partitionId, brokerUrl).isPresent());
+
+    // Different brokerUrls are stored independently
+    byte[] value1 = "value1".getBytes();
+    byte[] value2 = "value2".getBytes();
+    rocksDBStorageEngine.putGlobalRtDivState(partitionId, "broker1:9092", value1);
+    rocksDBStorageEngine.putGlobalRtDivState(partitionId, "broker2:9092", value2);
+    Assert.assertEquals(rocksDBStorageEngine.getGlobalRtDivState(partitionId, "broker1:9092").get(), value1);
+    Assert.assertEquals(rocksDBStorageEngine.getGlobalRtDivState(partitionId, "broker2:9092").get(), value2);
+
+    // Clean up
+    rocksDBStorageEngine.clearGlobalRtDivState(partitionId, "broker1:9092");
+    rocksDBStorageEngine.clearGlobalRtDivState(partitionId, "broker2:9092");
+  }
+
+  @Test
+  public void testIllegalPartitionIdInGlobalRtDivState() {
+    RocksDBStorageEngine rocksDBStorageEngine = (RocksDBStorageEngine) getTestStoreEngine();
+    String brokerUrl = "localhost:9092";
+    byte[] valueBytes = "value".getBytes();
+
+    // Negative partition id should throw
+    Assert.assertThrows(
+        IllegalArgumentException.class,
+        () -> rocksDBStorageEngine.putGlobalRtDivState(-1, brokerUrl, valueBytes));
+    Assert.assertThrows(IllegalArgumentException.class, () -> rocksDBStorageEngine.getGlobalRtDivState(-1, brokerUrl));
+    Assert
+        .assertThrows(IllegalArgumentException.class, () -> rocksDBStorageEngine.clearGlobalRtDivState(-1, brokerUrl));
+
+    // Metadata partition id should throw
+    Assert.assertThrows(
+        IllegalArgumentException.class,
+        () -> rocksDBStorageEngine.putGlobalRtDivState(METADATA_PARTITION_ID, brokerUrl, valueBytes));
+    Assert.assertThrows(
+        IllegalArgumentException.class,
+        () -> rocksDBStorageEngine.getGlobalRtDivState(METADATA_PARTITION_ID, brokerUrl));
+    Assert.assertThrows(
+        IllegalArgumentException.class,
+        () -> rocksDBStorageEngine.clearGlobalRtDivState(METADATA_PARTITION_ID, brokerUrl));
+  }
 }