@Nullable
  public static Integer getRealtimeSegmentPartitionId(String segmentName, String realtimeTableName,
      HelixManager helixManager, @Nullable String partitionColumn) {
    Integer partitionId = getPartitionIdFromRealtimeSegmentName(segmentName);
    if (partitionId != null) {
      return partitionId;
    }
    // Otherwise, retrieve the partition id from the segment zk metadata.
    SegmentZKMetadata segmentZKMetadata =
        ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), realtimeTableName, segmentName);
    Preconditions.checkState(segmentZKMetadata != null,
        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, realtimeTableName);
    return getRealtimeSegmentPartitionId(segmentZKMetadata, partitionColumn);
  }
@Test
  public void testGetRealtimeSegmentPartitionIdFromZkMetadata() {

    // mocks
    SegmentZKMetadata segmentZKMetadata = mock(SegmentZKMetadata.class);
    SegmentPartitionMetadata segmentPartitionMetadata = mock(SegmentPartitionMetadata.class);
    HashMap<String, ColumnPartitionMetadata> columnPartitionMetadataMap = new HashMap<>();
    HashSet<Integer> partitions = new HashSet<>();
    partitions.add(3);
    columnPartitionMetadataMap.put(PARTITION_COLUMN,
        new ColumnPartitionMetadata("modulo", 8, partitions, new HashMap<>()));

    when(segmentPartitionMetadata.getColumnPartitionMap()).thenReturn(columnPartitionMetadataMap);
    when(segmentZKMetadata.getPartitionMetadata()).thenReturn(segmentPartitionMetadata);

    HelixManager helixManager = mock(HelixManager.class);
    ZkHelixPropertyStore zkHelixPropertyStore = mock(ZkHelixPropertyStore.class);
    when(helixManager.getHelixPropertyStore()).thenReturn(zkHelixPropertyStore);

    // mock static ZKMetadataProvider.getSegmentZKMetadata
    try (MockedStatic<ZKMetadataProvider> zkMetadataProviderMockedStatic = Mockito.mockStatic(
        ZKMetadataProvider.class)) {
      when(ZKMetadataProvider.getSegmentZKMetadata(Mockito.any(ZkHelixPropertyStore.class), eq(TABLE_NAME_WITH_TYPE),
          eq(SEGMENT))).thenReturn(segmentZKMetadata);

      Integer partitionId =
          SegmentUtils.getRealtimeSegmentPartitionId(SEGMENT, TABLE_NAME_WITH_TYPE, helixManager, PARTITION_COLUMN);

      assertEquals(partitionId, 3);
    }
  }