Java Code Examples for org.apache.samza.system.SystemStreamMetadata#SystemStreamPartitionMetadata

@Test
public void testGetSSPMetadataEmptyPartition() {
  SystemStreamPartition ssp = new SystemStreamPartition(TEST_SYSTEM, VALID_TOPIC, new Partition(0));
  SystemStreamPartition otherSSP = new SystemStreamPartition(TEST_SYSTEM, "otherTopic", new Partition(1));
  TopicPartition topicPartition = new TopicPartition(VALID_TOPIC, 0);
  TopicPartition otherTopicPartition = new TopicPartition("otherTopic", 1);
  when(mockKafkaConsumer.beginningOffsets(ImmutableList.of(topicPartition, otherTopicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 1L));
  when(mockKafkaConsumer.endOffsets(ImmutableList.of(topicPartition, otherTopicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 11L));

  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> expected =
      ImmutableMap.of(ssp, new SystemStreamMetadata.SystemStreamPartitionMetadata("1", "10", "11"), otherSSP,
          new SystemStreamMetadata.SystemStreamPartitionMetadata(null, null, null));
  assertEquals(expected, kafkaSystemAdmin.getSSPMetadata(ImmutableSet.of(ssp, otherSSP)));
}
 

@Test
public void testGetSSPMetadataWithRetry() {
  SystemStreamPartition oneSSP = new SystemStreamPartition(TEST_SYSTEM, VALID_TOPIC, new Partition(0));
  SystemStreamPartition otherSSP = new SystemStreamPartition(TEST_SYSTEM, "otherTopic", new Partition(1));
  ImmutableSet<SystemStreamPartition> ssps = ImmutableSet.of(oneSSP, otherSSP);
  List<TopicPartition> topicPartitions = ssps.stream()
      .map(ssp -> new TopicPartition(ssp.getStream(), ssp.getPartition().getPartitionId()))
      .collect(Collectors.toList());
  Map<TopicPartition, Long> testBeginningOffsets =
      ImmutableMap.of(testTopicPartition0, KAFKA_BEGINNING_OFFSET_FOR_PARTITION0, testTopicPartition1,
          KAFKA_BEGINNING_OFFSET_FOR_PARTITION1);

  when(mockKafkaConsumer.beginningOffsets(topicPartitions)).thenThrow(new RuntimeException())
      .thenReturn(testBeginningOffsets);
  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> sspMetadata =
      kafkaSystemAdmin.getSSPMetadata(ssps, new ExponentialSleepStrategy(2,
          1, 1));

  assertEquals("metadata should return for 2 topics", sspMetadata.size(), 2);

  // retried twice because the first fails and the second succeeds
  Mockito.verify(mockKafkaConsumer, Mockito.times(2)).beginningOffsets(topicPartitions);
}
 

/**
 * This test is for cases, when calls to systemAdmin (e.g., KafkaSystemAdmin's) get-stream-metadata method return null.
 */
@Test
public void testGetStartingOffsetsWhenStreamMetadataIsNull() {
  final String taskName = "test-get-starting-offset-task";

  Set<SystemStreamPartition> ssps = IntStream.range(1, 2)
      .mapToObj(idx -> new SystemStreamPartition(TEST_SYSTEM, TEST_STREAM, new Partition(idx)))
      .collect(Collectors.toSet());
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = ssps.stream()
      .collect(Collectors.toMap(SystemStreamPartition::getPartition,
        x -> new SystemStreamMetadata.SystemStreamPartitionMetadata(null, "1", "2")));


  TaskSideInputHandler handler = new MockTaskSideInputHandlerBuilder(taskName, TaskMode.Active)
      .addStreamMetadata(Collections.singletonMap(new SystemStream(TEST_SYSTEM, TEST_STREAM),
          new SystemStreamMetadata(TEST_STREAM, partitionMetadata)))
      .addStore(TEST_STORE, ssps)
      .build();

  handler.init();

  ssps.forEach(ssp -> {
    String startingOffset = handler.getStartingOffset(
        new SystemStreamPartition(TEST_SYSTEM, TEST_STREAM, ssp.getPartition()));
    Assert.assertNull("Starting offset should be null", startingOffset);
  });
}
 

@Test
public void testIntegrationWithPartitioner() throws Exception {
  URL url = this.getClass().getResource("/partitioner");
  String whiteList = ".*";
  String blackList = ".*02";
  String groupPattern = "";
  String streamName = String.format(url.getPath());
  DirectoryPartitioner directoryPartitioner =
    new DirectoryPartitioner(whiteList, blackList, groupPattern, new HdfsFileSystemAdapter());
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> metadataMap = directoryPartitioner.getPartitionMetadataMap(streamName, null);
  Assert.assertEquals(1, metadataMap.size());
  Map<Partition, List<String>> descriporMap = directoryPartitioner.getPartitionDescriptor(streamName);
  Assert.assertEquals(1, descriporMap.values().size());
  Assert.assertTrue(descriporMap.get(new Partition(0)).get(0).endsWith("testfile01"));
}
 

@Test
public void testBroadcast() {
  SystemStreamMetadata metadata = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>();
  partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata);
  StreamMetadataCache metadataCache = mock(StreamMetadataCache.class);
  when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata);

  SystemStream systemStream = new SystemStream("test-system", "test-stream");
  Set<Integer> partitions = new HashSet<>();
  MessageCollector collector = mock(MessageCollector.class);
  doAnswer(invocation -> {
    OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0];
    partitions.add((Integer) envelope.getPartitionKey());
    assertEquals(envelope.getSystemStream(), systemStream);
    return null;
  }).when(collector).send(any());

  ControlMessageSender sender = new ControlMessageSender(metadataCache);
  WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0");
  SystemStreamPartition ssp = new SystemStreamPartition(systemStream, new Partition(0));
  sender.broadcastToOtherPartitions(watermark, ssp, collector);
  assertEquals(partitions.size(), 3);
}
 

@Test
public void testSend() {
  SystemStreamMetadata metadata = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>();
  partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata);
  StreamMetadataCache metadataCache = mock(StreamMetadataCache.class);
  when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata);

  SystemStream systemStream = new SystemStream("test-system", "test-stream");
  Set<Integer> partitions = new HashSet<>();
  MessageCollector collector = mock(MessageCollector.class);
  doAnswer(invocation -> {
    OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0];
    partitions.add((Integer) envelope.getPartitionKey());
    assertEquals(envelope.getSystemStream(), systemStream);
    return null;
  }).when(collector).send(any());

  ControlMessageSender sender = new ControlMessageSender(metadataCache);
  WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0");
  sender.send(watermark, systemStream, collector);
  assertEquals(partitions.size(), 1);
}
 

static SystemAdmin createSystemAdmin(Map<String, Integer> streamToPartitions) {

    return new SystemAdmin() {
      @Override
      public Map<SystemStreamPartition, String> getOffsetsAfter(Map<SystemStreamPartition, String> offsets) {
        return null;
      }

      @Override
      public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
        Map<String, SystemStreamMetadata> map = new HashMap<>();
        for (String stream : streamNames) {
          Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> m = new HashMap<>();
          for (int i = 0; i < streamToPartitions.get(stream); i++) {
            m.put(new Partition(i), new SystemStreamMetadata.SystemStreamPartitionMetadata("", "", ""));
          }
          map.put(stream, new SystemStreamMetadata(stream, m));
        }
        return map;
      }

      @Override
      public Integer offsetComparator(String offset1, String offset2) {
        return null;
      }
    };
  }
 

@Test
public void testGetStreamPartitionCounts() {
  SystemAdmin admin1 = mock(SystemAdmin.class);
  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  when(systemAdmins.getSystemAdmin(SYSTEM1)).thenReturn(admin1);

  Map<String, SystemStreamMetadata> map = new HashMap<>();
  SystemStreamMetadata meta1 = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitions = new HashMap<>();
  partitions.put(new Partition(0), null);
  when(meta1.getSystemStreamPartitionMetadata()).thenReturn(partitions);
  map.put(STREAM1, meta1);

  SystemStreamMetadata meta2 = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitions2 = new HashMap<>();
  partitions2.put(new Partition(0), null);
  partitions2.put(new Partition(1), null);
  when(meta2.getSystemStreamPartitionMetadata()).thenReturn(partitions2);
  map.put(STREAM2, meta2);

  when(admin1.getSystemStreamMetadata(anyObject())).thenReturn(map);

  Set<String> streams = new HashSet<>();
  streams.add(STREAM1);
  streams.add(STREAM2);
  StreamManager manager = new StreamManager(systemAdmins);
  Map<String, Integer> counts = manager.getStreamPartitionCounts(SYSTEM1, streams);

  assertTrue(counts.get(STREAM1).equals(1));
  assertTrue(counts.get(STREAM2).equals(2));
}
 

private void checkSideInputCaughtUp(SystemStreamPartition ssp, String offset, SystemStreamMetadata.OffsetType offsetType, boolean isEndOfStream) {

    if (isEndOfStream) {
      this.initialSideInputSSPMetadata.remove(ssp);
      this.sideInputsCaughtUp.countDown();
      LOG.info("Side input ssp {} has caught up to offset {} ({}).", ssp, offset, offsetType);
      return;
    }

    SystemStreamMetadata.SystemStreamPartitionMetadata sspMetadata = this.initialSideInputSSPMetadata.get(ssp);
    String offsetToCheck = sspMetadata == null ? null : sspMetadata.getOffset(offsetType);
    LOG.trace("Checking {} offset {} against {} for {}.", offsetType, offset, offsetToCheck, ssp);

    // Let's compare offset of the chosen message with offsetToCheck.
    Integer comparatorResult;
    if (offset == null || offsetToCheck == null) {
      comparatorResult = -1;
    } else {
      SystemAdmin systemAdmin = systemAdmins.getSystemAdmin(ssp.getSystem());
      comparatorResult = systemAdmin.offsetComparator(offset, offsetToCheck);
    }

    // The SSP is no longer lagging if the envelope's offset is greater than or equal to the
    // latest offset.
    if (comparatorResult != null && comparatorResult.intValue() >= 0) {

      LOG.info("Side input ssp {} has caught up to offset {} ({}).", ssp, offset, offsetType);
      // if its caught up, we remove the ssp from the map, and countDown the latch
      this.initialSideInputSSPMetadata.remove(ssp);
      this.sideInputsCaughtUp.countDown();
      return;
    }
  }
 

@Ignore("Integration Test")
@Test
public void testGetStreamMetadata() {
  EventHubSystemFactory eventHubSystemFactory = new EventHubSystemFactory();
  SystemAdmin eventHubSystemAdmin = eventHubSystemFactory.getAdmin(SYSTEM_NAME,
          MockEventHubConfigFactory.getEventHubConfig(EventHubSystemProducer.PartitioningMethod.EVENT_HUB_HASHING));
  Set<String> streams = new HashSet<>();
  streams.add(STREAM_NAME1);
  streams.add(STREAM_NAME2);
  Map<String, SystemStreamMetadata> metadataMap = eventHubSystemAdmin.getSystemStreamMetadata(streams);

  for (String stream : streams) {
    Assert.assertTrue(metadataMap.containsKey(stream));
    Assert.assertEquals(stream, metadataMap.get(stream).getStreamName());
    Assert.assertNotNull(metadataMap.get(stream).getSystemStreamPartitionMetadata());

    Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadataMap =
            metadataMap.get(stream).getSystemStreamPartitionMetadata();
    Assert.assertTrue(partitionMetadataMap.size() >= MIN_EVENTHUB_ENTITY_PARTITION);
    Assert.assertTrue(partitionMetadataMap.size() <= MAX_EVENTHUB_ENTITY_PARTITION);
    partitionMetadataMap.forEach((partition, metadata) -> {
      Assert.assertEquals(EventHubSystemConsumer.START_OF_STREAM, metadata.getOldestOffset());
      Assert.assertNotSame(EventHubSystemConsumer.END_OF_STREAM, metadata.getNewestOffset());
      String expectedUpcomingOffset = String.valueOf(Long.parseLong(metadata.getNewestOffset()) + 1);
      Assert.assertEquals(expectedUpcomingOffset, metadata.getUpcomingOffset());
    });
  }
}
 

@Test
public void testGetSSPMetadataZeroUpcomingOffset() {
  SystemStreamPartition ssp = new SystemStreamPartition(TEST_SYSTEM, VALID_TOPIC, new Partition(0));
  TopicPartition topicPartition = new TopicPartition(VALID_TOPIC, 0);
  when(mockKafkaConsumer.beginningOffsets(ImmutableList.of(topicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, -1L));
  when(mockKafkaConsumer.endOffsets(ImmutableList.of(topicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 0L));
  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> expected =
      ImmutableMap.of(ssp, new SystemStreamMetadata.SystemStreamPartitionMetadata("0", null, "0"));
  assertEquals(kafkaSystemAdmin.getSSPMetadata(ImmutableSet.of(ssp)), expected);
}
 

@Test
public void testGetSSPMetadataEmptyUpcomingOffset() {
  SystemStreamPartition ssp = new SystemStreamPartition(TEST_SYSTEM, VALID_TOPIC, new Partition(0));
  TopicPartition topicPartition = new TopicPartition(VALID_TOPIC, 0);
  when(mockKafkaConsumer.beginningOffsets(ImmutableList.of(topicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 0L));
  when(mockKafkaConsumer.endOffsets(ImmutableList.of(topicPartition))).thenReturn(ImmutableMap.of());
  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> expected =
      ImmutableMap.of(ssp, new SystemStreamMetadata.SystemStreamPartitionMetadata("0", null, null));
  assertEquals(kafkaSystemAdmin.getSSPMetadata(ImmutableSet.of(ssp)), expected);
}
 

@Test
public void testGetSSPMetadata() {
  SystemStreamPartition ssp = new SystemStreamPartition(TEST_SYSTEM, VALID_TOPIC, new Partition(0));
  SystemStreamPartition otherSSP = new SystemStreamPartition(TEST_SYSTEM, "otherTopic", new Partition(1));
  TopicPartition topicPartition = new TopicPartition(VALID_TOPIC, 0);
  TopicPartition otherTopicPartition = new TopicPartition("otherTopic", 1);
  when(mockKafkaConsumer.beginningOffsets(ImmutableList.of(topicPartition, otherTopicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 1L, otherTopicPartition, 2L));
  when(mockKafkaConsumer.endOffsets(ImmutableList.of(topicPartition, otherTopicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 11L, otherTopicPartition, 12L));
  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> expected =
      ImmutableMap.of(ssp, new SystemStreamMetadata.SystemStreamPartitionMetadata("1", "10", "11"), otherSSP,
          new SystemStreamMetadata.SystemStreamPartitionMetadata("2", "11", "12"));
  assertEquals(kafkaSystemAdmin.getSSPMetadata(ImmutableSet.of(ssp, otherSSP)), expected);
}
 

@Test
public void testGetSystemStreamMetaDataForTopicWithNoMessage() {
  // The topic with no messages will have beginningOffset = 0 and endOffset = 0
  when(mockKafkaConsumer.beginningOffsets(ImmutableList.of(testTopicPartition0, testTopicPartition1))).thenReturn(
      ImmutableMap.of(testTopicPartition0, 0L, testTopicPartition1, 0L));
  when(mockKafkaConsumer.endOffsets(ImmutableList.of(testTopicPartition0, testTopicPartition1))).thenReturn(
      ImmutableMap.of(testTopicPartition0, 0L, testTopicPartition1, 0L));

  Map<String, SystemStreamMetadata> metadataMap =
      kafkaSystemAdmin.getSystemStreamMetadata(ImmutableSet.of(VALID_TOPIC));
  assertEquals("metadata should return for 1 topic", metadataMap.size(), 1);

  // verify the metadata streamName
  assertEquals("the stream name should be " + VALID_TOPIC, metadataMap.get(VALID_TOPIC).getStreamName(), VALID_TOPIC);

  // verify the offset for each partition
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> systemStreamPartitionMetadata =
      metadataMap.get(VALID_TOPIC).getSystemStreamPartitionMetadata();
  assertEquals("there are 2 partitions", systemStreamPartitionMetadata.size(), 2);

  SystemStreamMetadata.SystemStreamPartitionMetadata partition0Metadata =
      systemStreamPartitionMetadata.get(new Partition(0));
  assertEquals("oldest offset for partition 0", partition0Metadata.getOldestOffset(), "0");
  assertEquals("upcoming offset for partition 0", partition0Metadata.getUpcomingOffset(), "0");
  assertEquals("newest offset is not set due to abnormal upcoming offset", partition0Metadata.getNewestOffset(),
      null);

  SystemStreamMetadata.SystemStreamPartitionMetadata partition1Metadata =
      systemStreamPartitionMetadata.get(new Partition(1));
  assertEquals("oldest offset for partition 1", partition1Metadata.getOldestOffset(), "0");
  assertEquals("upcoming offset for partition 1", partition1Metadata.getUpcomingOffset(), "0");
  assertEquals("newest offset is not set due to abnormal upcoming offset", partition1Metadata.getNewestOffset(),
      null);
}
 

@Test
public void testGetSystemStreamMetaDataWithValidTopic() {
  System.out.println("STARTING");
  Map<String, SystemStreamMetadata> metadataMap =
      kafkaSystemAdmin.getSystemStreamMetadata(ImmutableSet.of(VALID_TOPIC));

  // verify metadata size
  assertEquals("metadata should return for 1 topic", metadataMap.size(), 1);
  System.out.println("STARTING1");
  // verify the metadata streamName
  assertEquals("the stream name should be " + VALID_TOPIC, metadataMap.get(VALID_TOPIC).getStreamName(), VALID_TOPIC);
  System.out.println("STARTING2");
  // verify the offset for each partition
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> systemStreamPartitionMetadata =
      metadataMap.get(VALID_TOPIC).getSystemStreamPartitionMetadata();
  assertEquals("there are 2 partitions", systemStreamPartitionMetadata.size(), 2);
  System.out.println("STARTING3");
  SystemStreamMetadata.SystemStreamPartitionMetadata partition0Metadata =
      systemStreamPartitionMetadata.get(new Partition(0));
  assertEquals("oldest offset for partition 0", partition0Metadata.getOldestOffset(),
      KAFKA_BEGINNING_OFFSET_FOR_PARTITION0.toString());
  assertEquals("upcoming offset for partition 0", partition0Metadata.getUpcomingOffset(),
      KAFKA_END_OFFSET_FOR_PARTITION0.toString());
  assertEquals("newest offset for partition 0", partition0Metadata.getNewestOffset(),
      Long.toString(KAFKA_END_OFFSET_FOR_PARTITION0 - 1));
  System.out.println("STARTING4");
  SystemStreamMetadata.SystemStreamPartitionMetadata partition1Metadata =
      systemStreamPartitionMetadata.get(new Partition(1));
  assertEquals("oldest offset for partition 1", partition1Metadata.getOldestOffset(),
      KAFKA_BEGINNING_OFFSET_FOR_PARTITION1.toString());
  assertEquals("upcoming offset for partition 1", partition1Metadata.getUpcomingOffset(),
      KAFKA_END_OFFSET_FOR_PARTITION1.toString());
  assertEquals("newest offset for partition 1", partition1Metadata.getNewestOffset(),
      Long.toString(KAFKA_END_OFFSET_FOR_PARTITION1 - 1));
}
 

@Override
public Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> getSSPMetadata(
    Set<SystemStreamPartition> ssps) {
  return getSSPMetadata(ssps,
      new ExponentialSleepStrategy(DEFAULT_EXPONENTIAL_SLEEP_BACK_OFF_MULTIPLIER,
          DEFAULT_EXPONENTIAL_SLEEP_INITIAL_DELAY_MS, DEFAULT_EXPONENTIAL_SLEEP_MAX_DELAY_MS));
}
 

private SystemStreamMetadata constructSystemStreamMetadata(
    String streamName,
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> sspToMessagesForSystem) {

  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata =
      sspToMessagesForSystem
          .entrySet()
          .stream()
          .collect(Collectors.toMap(entry -> entry.getKey().getPartition(), entry -> {
            List<IncomingMessageEnvelope> messages = entry.getValue();
            Integer oldestOffset;
            Integer newestOffset;
            int upcomingOffset;

            if (messages.isEmpty()) {
              oldestOffset = null;
              newestOffset = null;
              upcomingOffset = 0;
            } else if (messages.get(messages.size() - 1).isEndOfStream()) {
              if (messages.size() > 1) {
                // don't count end of stream in offset indices
                oldestOffset = 0;
                newestOffset = messages.size() - 2;
                upcomingOffset = messages.size() - 1;
              } else {
                // end of stream is the only message, treat the same as empty
                oldestOffset = null;
                newestOffset = null;
                upcomingOffset = 0;
              }
            } else {
              // offsets correspond strictly to numeric indices
              oldestOffset = 0;
              newestOffset = messages.size() - 1;
              upcomingOffset = messages.size();
            }

            return new SystemStreamMetadata.SystemStreamPartitionMetadata(
                oldestOffset == null ? null : oldestOffset.toString(),
                newestOffset == null ? null : newestOffset.toString(),
                Integer.toString(upcomingOffset));
          }));

  return new SystemStreamMetadata(streamName, partitionMetadata);
}
 

@Test
public void testShouldAssembleMetadata() {
  Map<SystemStreamPartition, String> oldestOffsets = new ImmutableMap.Builder<SystemStreamPartition, String>()
      .put(new SystemStreamPartition(SYSTEM, "stream1", new Partition(0)), "o1")
      .put(new SystemStreamPartition(SYSTEM, "stream2", new Partition(0)), "o2")
      .put(new SystemStreamPartition(SYSTEM, "stream1", new Partition(1)), "o3")
      .put(new SystemStreamPartition(SYSTEM, "stream2", new Partition(1)), "o4")
      .build();

  Map<SystemStreamPartition, String> newestOffsets = new ImmutableMap.Builder<SystemStreamPartition, String>()
      .put(new SystemStreamPartition(SYSTEM, "stream1", new Partition(0)), "n1")
      .put(new SystemStreamPartition(SYSTEM, "stream2", new Partition(0)), "n2")
      .put(new SystemStreamPartition(SYSTEM, "stream1", new Partition(1)), "n3")
      .put(new SystemStreamPartition(SYSTEM, "stream2", new Partition(1)), "n4")
      .build();

  Map<SystemStreamPartition, String> upcomingOffsets = new ImmutableMap.Builder<SystemStreamPartition, String>()
      .put(new SystemStreamPartition(SYSTEM, "stream1", new Partition(0)), "u1")
      .put(new SystemStreamPartition(SYSTEM, "stream2", new Partition(0)), "u2")
      .put(new SystemStreamPartition(SYSTEM, "stream1", new Partition(1)), "u3")
      .put(new SystemStreamPartition(SYSTEM, "stream2", new Partition(1)), "u4")
      .build();

  Map<String, SystemStreamMetadata> metadata = assembleMetadata(oldestOffsets, newestOffsets, upcomingOffsets);
  assertNotNull(metadata);
  assertEquals(2, metadata.size());
  assertTrue(metadata.containsKey("stream1"));
  assertTrue(metadata.containsKey("stream2"));
  SystemStreamMetadata stream1Metadata = metadata.get("stream1");
  SystemStreamMetadata stream2Metadata = metadata.get("stream2");
  assertNotNull(stream1Metadata);
  assertNotNull(stream2Metadata);
  assertEquals("stream1", stream1Metadata.getStreamName());
  assertEquals("stream2", stream2Metadata.getStreamName());
  SystemStreamMetadata.SystemStreamPartitionMetadata expectedSystemStream1Partition0Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("o1", "n1", "u1");
  SystemStreamMetadata.SystemStreamPartitionMetadata expectedSystemStream1Partition1Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("o3", "n3", "u3");
  SystemStreamMetadata.SystemStreamPartitionMetadata expectedSystemStream2Partition0Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("o2", "n2", "u2");
  SystemStreamMetadata.SystemStreamPartitionMetadata expectedSystemStream2Partition1Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("o4", "n4", "u4");
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> stream1PartitionMetadata =
      stream1Metadata.getSystemStreamPartitionMetadata();
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> stream2PartitionMetadata =
      stream2Metadata.getSystemStreamPartitionMetadata();
  assertEquals(expectedSystemStream1Partition0Metadata, stream1PartitionMetadata.get(new Partition(0)));
  assertEquals(expectedSystemStream1Partition1Metadata, stream1PartitionMetadata.get(new Partition(1)));
  assertEquals(expectedSystemStream2Partition0Metadata, stream2PartitionMetadata.get(new Partition(0)));
  assertEquals(expectedSystemStream2Partition1Metadata, stream2PartitionMetadata.get(new Partition(1)));
}