ControlMessageSender { void send(ControlMessage message, SystemStream systemStream, MessageCollector collector) { int partitionCount = getPartitionCount(systemStream); int aggregatePartition = systemStream.getStream().hashCode() % partitionCount; LOG.debug(String.format("Send %s message from task %s to %s partition %s for aggregation", MessageType.of(message).name(), message.getTaskName(), systemStream, aggregatePartition)); OutgoingMessageEnvelope envelopeOut = new OutgoingMessageEnvelope(systemStream, aggregatePartition, null, message); collector.send(envelopeOut); } ControlMessageSender(StreamMetadataCache metadataCache);   }
@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); }