org.apache.samza.system.OutgoingMessageEnvelope Java Examples

@Test
public void testMultipleThreadWriteFlushInBoth() throws Exception {
  Thread t1 = writeFlushInThread(ome, azureBlobAvroWriter, 10);
  OutgoingMessageEnvelope ome2 = createOMEGenericRecord("TOPIC2");
  Thread t2 = writeFlushInThread(ome2, azureBlobAvroWriter, 10);

  t1.start();
  t2.start();
  t1.join(60000);
  t2.join(60000);

  verify(mockDataFileWriter, times(10)).appendEncoded(ByteBuffer.wrap(encodedRecord));
  verify(mockDataFileWriter, times(10)).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome2.getMessage())));
  verify(mockDataFileWriter, times(2)).flush();
  verify(mockAzureBlobOutputStream, times(20)).incrementNumberOfRecordsInBlob();
}
 

@Test
public void testEndToEndCompoundBooleanCheck() throws SamzaSqlValidatorException {

  int numMessages = 20;

  TestAvroSystemFactory.messages.clear();
  Map<String, String> staticConfigs = SamzaSqlTestConfig.fetchStaticConfigsWithFactories(numMessages);
  String sql1 = "Insert into testavro.outputTopic"
      + " select * from testavro.COMPLEX1 where id >= 0 and bool_value IS TRUE";
  List<String> sqlStmts = Arrays.asList(sql1);
  staticConfigs.put(SamzaSqlApplicationConfig.CFG_SQL_STMTS_JSON, JsonUtil.toJson(sqlStmts));

  Config config = new MapConfig(staticConfigs);
  new SamzaSqlValidator(config).validate(sqlStmts);

  runApplication(config);

  List<OutgoingMessageEnvelope> outMessages = new ArrayList<>(TestAvroSystemFactory.messages);
  Assert.assertEquals(numMessages / 2, outMessages.size());
}
 

@Override
public void window(MessageCollector collector, TaskCoordinator coordinator) {
  for (int i = 0; i < 100 && messagesSent < NewsfeedConfig.NUM_FOLLOW_EVENTS; i++, messagesSent++) {
    String follower = NewsfeedConfig.randomUser();
    String followee = NewsfeedConfig.randomUser();

    HashMap<String, Object> message = new HashMap<String, Object>();
    message.put("event", "follow");
    message.put("follower", follower);
    message.put("followee", followee);
    message.put("time", NewsfeedConfig.currentDateTime());
    collector.send(new OutgoingMessageEnvelope(NewsfeedConfig.FOLLOWS_STREAM, followee, null, message));
  }

  if (messagesSent % 100000 == 0) {
    log.info("Generated " + messagesSent + " follow events");
  }

  if (messagesSent == NewsfeedConfig.NUM_FOLLOW_EVENTS) {
    log.info("Finished generating random follower graph");
    coordinator.shutdown(RequestScope.CURRENT_TASK);
  }
}
 

protected EventData createEventData(String streamId, OutgoingMessageEnvelope envelope) {
  Optional<Interceptor> interceptor = Optional.ofNullable(interceptors.getOrDefault(streamId, null));
  byte[] eventValue = (byte[]) envelope.getMessage();
  if (interceptor.isPresent()) {
    eventValue = interceptor.get().intercept(eventValue);
  }

  EventData eventData = new EventDataImpl(eventValue);

  eventData.getProperties().put(PRODUCE_TIMESTAMP, Long.toString(System.currentTimeMillis()));

  if (config.getSendKeyInEventProperties(systemName)) {
    String keyValue = "";
    if (envelope.getKey() != null) {
      keyValue = (envelope.getKey() instanceof byte[]) ? new String((byte[]) envelope.getKey())
          : envelope.getKey().toString();
    }
    eventData.getProperties().put(KEY, keyValue);
  }
  return eventData;
}
 

@Test
public void testEndToEndWithBooleanCheck() throws Exception {
  int numMessages = 20;

  TestAvroSystemFactory.messages.clear();
  Map<String, String> staticConfigs = SamzaSqlTestConfig.fetchStaticConfigsWithFactories(numMessages);
  String sql1 = "Insert into testavro.outputTopic"
      + " select * from testavro.COMPLEX1 where bool_value IS TRUE";
  List<String> sqlStmts = Arrays.asList(sql1);
  staticConfigs.put(SamzaSqlApplicationConfig.CFG_SQL_STMTS_JSON, JsonUtil.toJson(sqlStmts));

  Config config = new MapConfig(staticConfigs);
  new SamzaSqlValidator(config).validate(sqlStmts);

  runApplication(config);

  List<OutgoingMessageEnvelope> outMessages = new ArrayList<>(TestAvroSystemFactory.messages);
  Assert.assertEquals(numMessages / 2, outMessages.size());
}
 

private void processWatermark(long timeStamp, MessageCollector collector) {
    if (timeStamp > lastWatermark + watermarkEvictionPeriod) {
        for (SlicingWindowOperator<Value> slicingWindowOperator : this.slicingWindowOperatorMap.values()) {
            List<AggregateWindow> aggregates = slicingWindowOperator.processWatermark(timeStamp);
            for (AggregateWindow<Value> aggregateWindow : aggregates) {
                if (aggregateWindow.hasValue()) {
                    System.out.println(aggregateWindow);
                    for (Value aggValue : aggregateWindow.getAggValues()) {
                        collector.send(new OutgoingMessageEnvelope(outputStream, aggValue));
                    }

                }
            }
        }
        lastWatermark = timeStamp;
    }
}
 

private StreamApplicationDescriptorImpl getKeyedSessionWindowStreamGraph(AccumulationMode mode, Duration duration) throws IOException {
  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    appDesc.getInputStream(inputDescriptor)
        .window(Windows.keyedSessionWindow(KV::getKey, duration, new IntegerSerde(), kvSerde)
            .setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

@Test
public void testEndToEndSubQuery() throws Exception {
  int numMessages = 20;
  TestAvroSystemFactory.messages.clear();
  Map<String, String> staticConfigs = SamzaSqlTestConfig.fetchStaticConfigsWithFactories(numMessages);
  String sql1 =
      "Insert into testavro.outputTopic(id, bool_value) select Flatten(a) as id, true as bool_value"
          + " from (select MyTestArray(id) a from testavro.SIMPLE1)";
  List<String> sqlStmts = Collections.singletonList(sql1);
  staticConfigs.put(SamzaSqlApplicationConfig.CFG_SQL_STMTS_JSON, JsonUtil.toJson(sqlStmts));

  Config config = new MapConfig(staticConfigs);
  new SamzaSqlValidator(config).validate(sqlStmts);

  runApplication(config);

  List<OutgoingMessageEnvelope> outMessages = new ArrayList<>(TestAvroSystemFactory.messages);

  int expectedMessages = 0;
  // Flatten de-normalizes the data. So there is separate record for each entry in the array.
  for (int index = 1; index < numMessages; index++) {
    expectedMessages = expectedMessages + Math.max(1, index);
  }
  Assert.assertEquals(expectedMessages, outMessages.size());
}
 

@Override
public void window(MessageCollector collector,
  TaskCoordinator coordinator) {

  KeyValueIterator<String, String> entries = store.all();
  while (entries.hasNext()) {                                        // c
    Entry<String, String> entry = entries.next();
    String key = entry.getKey();
    String value = entry.getValue();
    if (isTimestampKey(key) && Cart.isAbandoned(value)) {            // d
      String shopper = extractShopper(key);
      String cart = store.get(asCartKey(shopper));
      
      AbandonedCartEvent event =
        new AbandonedCartEvent(shopper, cart);
      collector.send(new OutgoingMessageEnvelope(
        new SystemStream("kafka", "derived-events-ch04"), event));    // e
      
      resetShopper(shopper);
    }
  }
}
 

private Thread writeFlushInThread(OutgoingMessageEnvelope ome, AzureBlobAvroWriter azureBlobAvroWriter,
    int numberOfSends) {
  Thread t = new Thread() {
    @Override
    public void run() {
      try {
        for (int i = 0; i < numberOfSends; i++) {
          azureBlobAvroWriter.write(ome);
        }
        azureBlobAvroWriter.flush();
      } catch (IOException e) {
        throw new SamzaException(e);
      }
    }
  };
  return t;
}
 

private Thread writeInThread(OutgoingMessageEnvelope ome, AzureBlobAvroWriter azureBlobAvroWriter,
    int numberOfSends) {
  Thread t = new Thread() {
    @Override
    public void run() {
      try {
        for (int i = 0; i < numberOfSends; i++) {
          azureBlobAvroWriter.write(ome);
        }
      } catch (IOException e) {
        throw new SamzaException(e);
      }
    }
  };
  return t;
}
 

private StreamApplicationDescriptorImpl getTumblingWindowStreamGraph(AccumulationMode mode,
    Duration duration, Trigger<KV<Integer, Integer>> earlyTrigger) throws IOException {
  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    appDesc.getInputStream(inputDescriptor)
        .window(Windows.tumblingWindow(duration, kvSerde).setEarlyTrigger(earlyTrigger)
            .setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

private StreamApplicationDescriptorImpl getKeyedTumblingWindowStreamGraph(AccumulationMode mode,
    Duration duration, Trigger<KV<Integer, Integer>> earlyTrigger) throws IOException {

  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    appDesc.getInputStream(inputDescriptor)
        .window(Windows.keyedTumblingWindow(KV::getKey, duration, new IntegerSerde(), kvSerde)
            .setEarlyTrigger(earlyTrigger).setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

private StreamApplicationDescriptorImpl getTestJoinStreamGraph(TestJoinFunction joinFn) throws IOException {
  Map<String, String> mapConfig = new HashMap<>();
  mapConfig.put("job.name", "jobName");
  mapConfig.put("job.id", "jobId");
  StreamTestUtils.addStreamConfigs(mapConfig, "inStream", "insystem", "instream");
  StreamTestUtils.addStreamConfigs(mapConfig, "inStream2", "insystem", "instream2");
  Config config = new MapConfig(mapConfig);

  return new StreamApplicationDescriptorImpl(appDesc -> {
    IntegerSerde integerSerde = new IntegerSerde();
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(integerSerde, integerSerde);
    GenericSystemDescriptor sd = new GenericSystemDescriptor("insystem", "mockFactoryClassName");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor1 = sd.getInputDescriptor("inStream", kvSerde);
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor2 = sd.getInputDescriptor("inStream2", kvSerde);

    MessageStream<KV<Integer, Integer>> inStream = appDesc.getInputStream(inputDescriptor1);
    MessageStream<KV<Integer, Integer>> inStream2 = appDesc.getInputStream(inputDescriptor2);

    inStream
        .join(inStream2, joinFn, integerSerde, kvSerde, kvSerde, JOIN_TTL, "j1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  }, config);
}
 

/**
 * Sends a specified message envelope from a specified Samza source.

 * @param source String representing the source of the message.
 * @param envelope Aggregate object representing the serialized message to send from the source.
 */
@Override
public void send(String source, OutgoingMessageEnvelope envelope) {
  Object key = envelope.getKey();
  Object message = envelope.getMessage();

  Object partitionKey;
  // We use the partition key from message if available, if not fallback to message key or use message as partition
  // key as the final resort.
  if (envelope.getPartitionKey() != null) {
    partitionKey = envelope.getPartitionKey();
  } else if (key != null) {
    partitionKey = key;
  } else {
    partitionKey = message;
  }

  Preconditions.checkNotNull(partitionKey, "Failed to compute partition key for the message: " + envelope);

  int partition =
      Math.abs(hashCode(partitionKey)) % memoryManager.getPartitionCountForSystemStream(envelope.getSystemStream());

  SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partition));
  memoryManager.put(ssp, key, message);
}
 

/**
 * Serialize and send a coordinator stream message.
 *
 * @param message
 *          The message to send.
 */
public void send(CoordinatorStreamMessage message) {
  log.debug("Sending {}", message);
  try {
    String source = message.getSource();
    byte[] key = keySerde.toBytes(Arrays.asList(message.getKeyArray()));
    byte[] value = null;
    if (!message.isDelete()) {
      value = messageSerde.toBytes(message.getMessageMap());
    }
    OutgoingMessageEnvelope envelope = new OutgoingMessageEnvelope(systemStream, Integer.valueOf(0), key, value);
    systemProducer.send(source, envelope);
  } catch (Exception e) {
    throw new SamzaException(e);
  }
}
 

@Test
public void testMultipleThreadWriteFlushFinallyClose() throws Exception {
  Thread t1 = writeFlushInThread(ome, azureBlobAvroWriter, 10);
  OutgoingMessageEnvelope ome2 = createOMEGenericRecord("TOPIC2");
  Thread t2 = writeFlushInThread(ome2, azureBlobAvroWriter, 10);

  t1.start();
  t2.start();
  t1.join(60000);
  t2.join(60000);
  azureBlobAvroWriter.close();

  verify(mockDataFileWriter, times(10)).appendEncoded(ByteBuffer.wrap(encodedRecord));
  verify(mockDataFileWriter, times(10)).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome2.getMessage())));
  verify(mockDataFileWriter, times(2)).flush();
  verify(mockDataFileWriter).close();
  verify(mockAzureBlobOutputStream, times(20)).incrementNumberOfRecordsInBlob();
}
 

private StreamApplicationDescriptorImpl getAggregateTumblingWindowStreamGraph(AccumulationMode mode, Duration timeDuration,
      Trigger<IntegerEnvelope> earlyTrigger) throws IOException {
  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    MessageStream<KV<Integer, Integer>> integers = appDesc.getInputStream(inputDescriptor);

    integers
        .map(new KVMapFunction())
        .window(Windows.<IntegerEnvelope, Integer>tumblingWindow(timeDuration, () -> 0, (m, c) -> c + 1, new IntegerSerde())
            .setEarlyTrigger(earlyTrigger)
            .setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

@Override
public IndexRequest getIndexRequest(OutgoingMessageEnvelope envelope) {
  IndexRequest indexRequest = getRequest(envelope);

  getId(envelope).ifPresent(indexRequest::id);
  getRoutingKey(envelope).ifPresent(indexRequest::routing);
  getVersion(envelope).ifPresent(indexRequest::version);
  getVersionType(envelope).ifPresent(indexRequest::versionType);

  setSource(envelope, indexRequest);

  return indexRequest;
}
 

@Override
public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator)
    throws Exception {
  Integer obj = (Integer) envelope.getMessage();
  collector.send(new OutgoingMessageEnvelope(new SystemStream("test", "output"),
      envelope.getKey(), envelope.getKey(), obj * multiplier));
}
 

@Test
public void testSecondPublishWithExceptionInSecondMessage() {
  // Across two successive run() invocations two messages should be published if stop events are added
  this.diagnosticsManager.start();
  DiagnosticsExceptionEvent diagnosticsExceptionEvent = new DiagnosticsExceptionEvent(System.currentTimeMillis(), new RuntimeException("exception"), new HashMap());
  this.diagnosticsManager.addExceptionEvent(diagnosticsExceptionEvent);
  this.diagnosticsManager.start();

  Assert.assertEquals("Two messages should have been published", 2, mockSystemProducer.getEnvelopeList().size());

  // Validate the first message
  OutgoingMessageEnvelope outgoingMessageEnvelope = mockSystemProducer.getEnvelopeList().get(0);
  validateMetricsHeader(outgoingMessageEnvelope);
  validateOutgoingMessageEnvelope(outgoingMessageEnvelope);

  // Validate the second message's header
  outgoingMessageEnvelope = mockSystemProducer.getEnvelopeList().get(1);
  validateMetricsHeader(outgoingMessageEnvelope);

  // Validate the second message's body (should be all empty except for the processor-stop-event)
  MetricsSnapshot metricsSnapshot =
      new MetricsSnapshotSerdeV2().fromBytes((byte[]) outgoingMessageEnvelope.getMessage());
  DiagnosticsStreamMessage diagnosticsStreamMessage =
      DiagnosticsStreamMessage.convertToDiagnosticsStreamMessage(metricsSnapshot);

  Assert.assertNull(diagnosticsStreamMessage.getContainerMb());
  Assert.assertEquals(Arrays.asList(diagnosticsExceptionEvent), diagnosticsStreamMessage.getExceptionEvents());
  Assert.assertNull(diagnosticsStreamMessage.getProcessorStopEvents());
  Assert.assertNull(diagnosticsStreamMessage.getContainerModels());
  Assert.assertNull(diagnosticsStreamMessage.getContainerNumCores());
  Assert.assertNull(diagnosticsStreamMessage.getNumPersistentStores());
}
 

void broadcastToOtherPartitions(ControlMessage message, SystemStreamPartition ssp, MessageCollector collector) {
  SystemStream systemStream = ssp.getSystemStream();
  int partitionCount = getPartitionCount(systemStream);
  int currentPartition = ssp.getPartition().getPartitionId();
  for (int i = 0; i < partitionCount; i++) {
    if (i != currentPartition) {
      OutgoingMessageEnvelope envelopeOut = new OutgoingMessageEnvelope(systemStream, i, null, message);
      collector.send(envelopeOut);
    }
  }
}
 

/**
 * Test keys of type byte[] goes to the same partition if they have the same contents.
 */
@Test
public void testPartition() {
  doReturn(1000).when(inMemoryManager).getPartitionCountForSystemStream(any());
  doAnswer(new Answer<Void>() {
    int partitionOfFirstMessage = -1;
    int partitionOfSecondMessage = -2;

    @Override
    public Void answer(InvocationOnMock invocation) throws Throwable {
      SystemStreamPartition ssp = invocation.getArgumentAt(0, SystemStreamPartition.class);
      if (partitionOfFirstMessage == -1) {
        partitionOfFirstMessage = ssp.getPartition().getPartitionId();
      } else {
        partitionOfSecondMessage = ssp.getPartition().getPartitionId();
        Assert.assertEquals(partitionOfFirstMessage, partitionOfSecondMessage);
        testFinished = true;
      }
      return null;
    }
  }).when(inMemoryManager).put(any(), any(), any());

  byte[] key1 = new byte[]{1, 2, 3};
  byte[] key2 = new byte[]{1, 2, 3};
  SystemStream systemStream = new SystemStream("TestSystem", "TestStream");
  OutgoingMessageEnvelope outgoingMessageEnvelope1 = new OutgoingMessageEnvelope(systemStream, key1, null);
  OutgoingMessageEnvelope outgoingMessageEnvelope2 = new OutgoingMessageEnvelope(systemStream, key2, null);
  inMemorySystemProducer.send("TestSource", outgoingMessageEnvelope1);
  inMemorySystemProducer.send("TestSource", outgoingMessageEnvelope2);
  Assert.assertTrue(testFinished);
}
 

@Override
public void process(
    IncomingMessageEnvelope incomingMessageEnvelope,
    MessageCollector messageCollector,
    TaskCoordinator taskCoordinator) throws Exception {
  messageCollector.send(
      new OutgoingMessageEnvelope(
          new SystemStream(outputSystem, outputTopic),
          incomingMessageEnvelope.getMessage()));
  processedMessageCount++;
  if (processedMessageCount == expectedMessageCount) {
    taskCoordinator.shutdown(TaskCoordinator.RequestScope.ALL_TASKS_IN_CONTAINER);
  }
}
 

@Test
public void testCoordinatorStreamSystemProducer() {
  MockCoordinatorStreamSystemFactory.enableMockConsumerCache();
  String source = "source";
  SystemStream systemStream = new SystemStream("system", "stream");
  MockCoordinatorSystemProducer systemProducer = new MockCoordinatorSystemProducer(source);
  MockSystemAdmin systemAdmin = new MockSystemAdmin();
  CoordinatorStreamSystemProducer producer = new CoordinatorStreamSystemProducer(systemStream, systemProducer, systemAdmin);
  SetConfig setConfig1 = new SetConfig(source, "job.name", "my-job-name");
  SetConfig setConfig2 = new SetConfig(source, "job.id", "1234");
  Delete delete = new Delete(source, "job.name", SetConfig.TYPE);
  assertFalse(systemProducer.isRegistered());
  producer.register(source);
  assertTrue(systemProducer.isRegistered());
  assertFalse(systemProducer.isStarted());
  producer.start();
  assertTrue(systemProducer.isStarted());
  producer.send(setConfig1);
  producer.send(setConfig2);
  producer.send(delete);
  assertFalse(systemProducer.isStopped());
  producer.stop();
  assertTrue(systemProducer.isStopped());
  List<OutgoingMessageEnvelope> envelopes = systemProducer.getEnvelopes();
  OutgoingMessageEnvelope envelope0 = envelopes.get(0);
  OutgoingMessageEnvelope envelope1 = envelopes.get(1);
  OutgoingMessageEnvelope envelope2 = envelopes.get(2);
  TypeReference<Object[]> keyRef = new TypeReference<Object[]>() {
  };
  TypeReference<Map<String, Object>> msgRef = new TypeReference<Map<String, Object>>() {
  };
  assertEquals(3, envelopes.size());
  assertEquals(new CoordinatorStreamMessage(setConfig1), new CoordinatorStreamMessage(deserialize((byte[]) envelope0.getKey(), keyRef), deserialize((byte[]) envelope0.getMessage(), msgRef)));
  assertEquals(new CoordinatorStreamMessage(setConfig2), new CoordinatorStreamMessage(deserialize((byte[]) envelope1.getKey(), keyRef), deserialize((byte[]) envelope1.getMessage(), msgRef)));
  assertEquals(new CoordinatorStreamMessage(delete), new CoordinatorStreamMessage(deserialize((byte[]) envelope2.getKey(), keyRef), deserialize((byte[]) envelope2.getMessage(), msgRef)));
}
 

@Test
public void testSend() throws Exception {
  OutgoingMessageEnvelope envelope = mock(OutgoingMessageEnvelope.class);
  IndexRequest indexRequest = mock(IndexRequest.class);

  when(INDEX_REQUEST_FACTORY.getIndexRequest(envelope)).thenReturn(indexRequest);

  producer.send(SOURCE_ONE, envelope);

  verify(processorOne).add(indexRequest);
}
 

private String[] getFormattedRow(OutgoingMessageEnvelope row) {
  String[] formattedRow = new String[1];
  String outputFormat = environmentVariableHandler.getEnvironmentVariable(SAMZA_SQL_OUTPUT);
  if (outputFormat == null || !outputFormat.equalsIgnoreCase(MessageFormat.PRETTY.toString())) {
    formattedRow[0] = getCompressedFormat(row);
  } else {
    formattedRow[0] = getPrettyFormat(row);
  }
  return formattedRow;
}
 

@Test
public void testSecondPublishWithProcessorStopInSecondMessage() {
  // Across two successive run() invocations two messages should be published if stop events are added
  this.diagnosticsManager.start();
  this.diagnosticsManager.addProcessorStopEvent("0", executionEnvContainerId, hostname, 102);
  this.diagnosticsManager.start();

  Assert.assertEquals("Two messages should have been published", 2, mockSystemProducer.getEnvelopeList().size());

  // Validate the first message
  OutgoingMessageEnvelope outgoingMessageEnvelope = mockSystemProducer.getEnvelopeList().get(0);
  validateMetricsHeader(outgoingMessageEnvelope);
  validateOutgoingMessageEnvelope(outgoingMessageEnvelope);

  // Validate the second message's header
  outgoingMessageEnvelope = mockSystemProducer.getEnvelopeList().get(1);
  validateMetricsHeader(outgoingMessageEnvelope);

  // Validate the second message's body (should be all empty except for the processor-stop-event)
  MetricsSnapshot metricsSnapshot =
      new MetricsSnapshotSerdeV2().fromBytes((byte[]) outgoingMessageEnvelope.getMessage());
  DiagnosticsStreamMessage diagnosticsStreamMessage =
      DiagnosticsStreamMessage.convertToDiagnosticsStreamMessage(metricsSnapshot);

  Assert.assertNull(diagnosticsStreamMessage.getContainerMb());
  Assert.assertNull(diagnosticsStreamMessage.getExceptionEvents());
  Assert.assertEquals(diagnosticsStreamMessage.getProcessorStopEvents(),
      Arrays.asList(new ProcessorStopEvent("0", executionEnvContainerId, hostname, 102)));
  Assert.assertNull(diagnosticsStreamMessage.getContainerModels());
  Assert.assertNull(diagnosticsStreamMessage.getContainerNumCores());
  Assert.assertNull(diagnosticsStreamMessage.getNumPersistentStores());
}
 

@Override
public void describe(StreamApplicationDescriptor appDescriptor) {
  KafkaSystemDescriptor ksd = new KafkaSystemDescriptor("test");
  KafkaInputDescriptor<KV<String, PageView>> isd =
      ksd.getInputDescriptor("PageView", KVSerde.of(new NoOpSerde<>(), new NoOpSerde<>()));
  MessageStream<KV<String, TestData.PageView>> inputStream = appDescriptor.getInputStream(isd);
  inputStream
      .map(KV::getValue)
      .partitionBy(PageView::getMemberId, pv -> pv, KVSerde.of(new IntegerSerde(), new JsonSerdeV2<>(PageView.class)), "p1")
      .sink((m, collector, coordinator) ->
          collector.send(new OutgoingMessageEnvelope(new SystemStream("test", "Output"), m.getKey(), m.getKey(), m)));
}
 

@Override
public List<String[]> consumeQueryResult(ExecutionContext context, int startRow, int endRow) {
  List<String[]> results = new ArrayList<>();
  for (OutgoingMessageEnvelope row : outputData.consume(startRow, endRow)) {
    results.add(getFormattedRow(row));
  }
  return results;
}