MessageStreamImpl implements MessageStream<M> { @Override public <K, V> MessageStream<KV<K, V>> partitionBy(MapFunction<? super M, ? extends K> keyExtractor, MapFunction<? super M, ? extends V> valueExtractor, KVSerde<K, V> serde, String userDefinedId) { String opId = this.streamAppDesc.getNextOpId(OpCode.PARTITION_BY, userDefinedId); IntermediateMessageStreamImpl<KV<K, V>> intermediateStream = this.streamAppDesc.getIntermediateStream(opId, serde, false); if (!intermediateStream.isKeyed()) { throw new SamzaException("partitionBy can not be used with a default serde that is not a KVSerde."); } PartitionByOperatorSpec<M, K, V> partitionByOperatorSpec = OperatorSpecs.createPartitionByOperatorSpec( intermediateStream.getOutputStream(), keyExtractor, valueExtractor, opId); this.operatorSpec.registerNextOperatorSpec(partitionByOperatorSpec); return intermediateStream; } MessageStreamImpl(StreamApplicationDescriptorImpl streamAppDesc, OperatorSpec<?, M> operatorSpec); @Override MessageStream<TM> map(MapFunction<? super M, ? extends TM> mapFn); @Override MessageStream<M> filter(FilterFunction<? super M> filterFn); @Override MessageStream<TM> flatMap(FlatMapFunction<? super M, ? extends TM> flatMapFn); @Override MessageStream<OM> flatMapAsync(AsyncFlatMapFunction<? super M, ? extends OM> flatMapFn); @Override void sink(SinkFunction<? super M> sinkFn); @Override MessageStream<M> sendTo(OutputStream<M> outputStream); @Override MessageStream<WindowPane<K, WV>> window(Window<M, K, WV> window, String userDefinedId); @Override MessageStream<JM> join(MessageStream<OM> otherStream,
      JoinFunction<? extends K, ? super M, ? super OM, ? extends JM> joinFn,
      Serde<K> keySerde, Serde<M> messageSerde, Serde<OM> otherMessageSerde,
      Duration ttl, String userDefinedId); @Override MessageStream<JM> join(Table<R> table,
      StreamTableJoinFunction<? extends K, ? super M, ? super R, ? extends JM> joinFn, Object ... args); @Override MessageStream<M> merge(Collection<? extends MessageStream<? extends M>> otherStreams); @Override MessageStream<KV<K, V>> partitionBy(MapFunction<? super M, ? extends K> keyExtractor,
      MapFunction<? super M, ? extends V> valueExtractor, KVSerde<K, V> serde, String userDefinedId); @Override MessageStream<KV<K, V>> sendTo(Table<KV<K, V>> table, Object ... args); @Override MessageStream<M> broadcast(Serde<M> serde, String userDefinedId); @VisibleForTesting OperatorSpec<?, M> getOperatorSpec();  }
@Test public void testPartitionBy() throws IOException { StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class); OperatorSpec mockOpSpec = mock(OperatorSpec.class); String mockOpName = "mockName"; when(mockGraph.getNextOpId(anyObject(), anyObject())).thenReturn(mockOpName); OutputStreamImpl mockOutputStreamImpl = mock(OutputStreamImpl.class); KVSerde mockKVSerde = mock(KVSerde.class); IntermediateMessageStreamImpl mockIntermediateStream = mock(IntermediateMessageStreamImpl.class); when(mockGraph.getIntermediateStream(eq(mockOpName), eq(mockKVSerde), eq(false))) .thenReturn(mockIntermediateStream); when(mockIntermediateStream.getOutputStream()) .thenReturn(mockOutputStreamImpl); when(mockIntermediateStream.isKeyed()).thenReturn(true); MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph, mockOpSpec); MapFunction mockKeyFunction = mock(MapFunction.class); MapFunction mockValueFunction = mock(MapFunction.class); inputStream.partitionBy(mockKeyFunction, mockValueFunction, mockKVSerde, "p1"); ArgumentCaptor<OperatorSpec> registeredOpCaptor = ArgumentCaptor.forClass(OperatorSpec.class); verify(mockOpSpec).registerNextOperatorSpec(registeredOpCaptor.capture()); OperatorSpec<?, TestMessageEnvelope> registeredOpSpec = registeredOpCaptor.getValue(); assertTrue(registeredOpSpec instanceof PartitionByOperatorSpec); assertEquals(OpCode.PARTITION_BY, registeredOpSpec.getOpCode()); assertEquals(mockOutputStreamImpl, ((PartitionByOperatorSpec) registeredOpSpec).getOutputStream()); assertEquals(mockKeyFunction, ((PartitionByOperatorSpec) registeredOpSpec).getKeyFunction()); assertEquals(mockValueFunction, ((PartitionByOperatorSpec) registeredOpSpec).getValueFunction()); }