org.mockito.AdditionalMatchers Java Examples

@SuppressWarnings("unchecked")
@Test
public void testMultipleSuccess() throws IOException {
    final Set<FlowFile> flowFiles = new HashSet<>();
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));

    when(mockLease.complete()).thenReturn(createAllSuccessPublishResult(flowFiles, 1));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_1_0.REL_SUCCESS, 3);

    verify(mockLease, times(3)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();
}
 

private void testRetryNeverSucceeds(final int maxAttemptCount, final int maxAttemptIntervalIncrease) {
    when(supplier.get()).thenThrow(new RuntimeException("something went wrong"));

    final Optional result = retryWithExponentialBackoff.retry(
            supplier,
            maxAttemptCount,
            maxAttemptIntervalIncrease,
            reportFailOnce
    );

    assertEquals(Optional.empty(), result);
    verify(supplier, times(maxAttemptCount)).get();
    verify(reportFailOnce, times(maxAttemptCount)).accept(any(Exception.class), anyInt());
    verify(sleep, times(maxAttemptCount - 1))
            .apply(AdditionalMatchers.leq((1L << maxAttemptIntervalIncrease) * 1000));
    verify(sleep, never())
            .apply(AdditionalMatchers.gt((1L << maxAttemptIntervalIncrease) * 1000));
}
 

@Test
public void testMultipleSuccess() throws IOException {
    final Set<FlowFile> flowFiles = new HashSet<>();
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));

    when(mockLease.complete()).thenReturn(createAllSuccessPublishResult(flowFiles, 1));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_2_0.REL_SUCCESS, 3);

    verify(mockLease, times(3)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();
}
 

@Test
public void testSliderSingleClickDifferentPosition_ListenerShouldBeCalled() {
  // Lays out slider.
  helper.addContentView(activity);

  // Click pressed.
  touchSliderAtValue(slider, SLIDER_VALUE_FROM + SLIDER_VALUE_RANGE / 2, MotionEvent.ACTION_DOWN);
  // Click released.
  touchSliderAtValue(slider, SLIDER_VALUE_FROM + SLIDER_VALUE_RANGE / 2, MotionEvent.ACTION_UP);
  // Listener should be called once.
  verify(mockOnChangeListener, times(1))
      .onValueChange(
          eq(slider),
          AdditionalMatchers.eq(SLIDER_VALUE_FROM + SLIDER_VALUE_RANGE / 2, FLOAT_ERROR),
          eq(true));
}
 

@Test
public void testGetAllModules_multi() {
  Mockito.when(server.getModules()).thenReturn(new IModule[] {module1});
  Mockito.when(server.getChildModules(any(IModule[].class), any(IProgressMonitor.class)))
      .thenReturn(new IModule[0]);

  Mockito.when(server.getChildModules(AdditionalMatchers.aryEq(new IModule[] {module1}),
      any(IProgressMonitor.class))).thenReturn(new IModule[] {module2a, module2b});

  Mockito.when(server.getChildModules(AdditionalMatchers.aryEq(new IModule[] {module1, module2b}),
      any(IProgressMonitor.class))).thenReturn(new IModule[] {module3});

  IModule[] result = ModuleUtils.getAllModules(server);
  Assert.assertNotNull(result);
  Assert.assertEquals(4, result.length);
  Assert.assertThat(result,
      IsArrayContainingInOrder.arrayContaining(module1, module2a, module2b, module3));
}
 

@SuppressWarnings("unchecked")
@Test
public void testSingleSuccess() throws IOException {
    final MockFlowFile flowFile = runner.enqueue("John Doe, 48");

    when(mockLease.complete()).thenReturn(createAllSuccessPublishResult(flowFile, 1));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_1_0.REL_SUCCESS, 1);

    verify(mockLease, times(1)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();
}
 

@Test
public void testNoRecordsInFlowFile() throws IOException {
    final List<FlowFile> flowFiles = new ArrayList<>();
    flowFiles.add(runner.enqueue(new byte[0]));

    final Map<FlowFile, Integer> msgCounts = new HashMap<>();
    msgCounts.put(flowFiles.get(0), 0);

    final PublishResult result = createPublishResult(msgCounts, new HashSet<>(flowFiles), Collections.emptyMap());

    when(mockLease.complete()).thenReturn(result);

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_0_10.REL_SUCCESS, 1);

    verify(mockLease, times(1)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();

    final MockFlowFile mff = runner.getFlowFilesForRelationship(PublishKafkaRecord_0_10.REL_SUCCESS).get(0);
    mff.assertAttributeEquals("msg.count", "0");
}
 

@Test
public void testMultipleFailures() throws IOException {
    final Set<FlowFile> flowFiles = new HashSet<>();
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));

    when(mockLease.complete()).thenReturn(createFailurePublishResult(flowFiles));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_2_0.REL_FAILURE, 3);

    verify(mockLease, times(3)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(1)).close();
}
 

@Test
public void playerSpecifcTest() {
  final String command1 = "/player1 test";
  final String command2 = "/player2 test2";

  final TabCompleteEvent event1p1 = generateTabCompleteEvent(connection1, command1);
  final TabCompleteEvent event1p2 = generateTabCompleteEvent(connection2, command1);
  final TabCompleteEvent event2p1 = generateTabCompleteEvent(connection1, command2);
  final TabCompleteEvent event2p2 = generateTabCompleteEvent(connection2, command2);

  listener.onTabComplete(event1p1);
  listener.onTabComplete(event1p2);
  listener.onTabComplete(event2p1);
  listener.onTabComplete(event2p2);

  Mockito.verify(commandPlayer1, Mockito.times(2))
      .tabComplete(Mockito.any(), AdditionalMatchers.aryEq(new String[] {"test"}));
  Mockito.verify(commandPlayer2, Mockito.times(2))
      .tabComplete(Mockito.any(), AdditionalMatchers.aryEq(new String[] {"test2"}));
  assertEquals(Arrays.asList("test"), event1p1.getSuggestions());
  assertEquals(Collections.emptyList(), event1p2.getSuggestions());
  assertEquals(Collections.emptyList(), event2p1.getSuggestions());
  assertEquals(Arrays.asList("test2"), event2p2.getSuggestions());
}
 

@Test
public void testMultipleFailures() throws IOException {
    final Set<FlowFile> flowFiles = new HashSet<>();
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));

    when(mockLease.complete()).thenReturn(createFailurePublishResult(flowFiles));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_0_11.REL_FAILURE, 3);

    verify(mockLease, times(3)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(1)).close();
}
 

@Test
public void testWriteKeyStoreTruncateFailure() throws IOException, GeneralSecurityException {
    setUnlimitedCrypto(false);
    String testPassword = "testPassword";
    String truncatedPassword = testPassword.substring(0, 7);
    IOException ioException = new IOException(TlsHelper.ILLEGAL_KEY_SIZE);
    IOException ioException2 = new IOException(TlsHelper.ILLEGAL_KEY_SIZE);
    doThrow(ioException).when(keyStoreSpi).engineStore(eq(tmpFileOutputStream), AdditionalMatchers.aryEq(testPassword.toCharArray()));
    doThrow(ioException2).when(keyStoreSpi).engineStore(eq(tmpFileOutputStream), AdditionalMatchers.aryEq(truncatedPassword.toCharArray()));
    try {
        TlsHelper.writeKeyStore(keyStore, outputStreamFactory, file, testPassword, true);
        fail("Expected " + ioException2);
    } catch (IOException e) {
        assertEquals(ioException2, e);
    }
}
 

@Test
public void testNoRecordsInFlowFile() throws IOException {
    final List<FlowFile> flowFiles = new ArrayList<>();
    flowFiles.add(runner.enqueue(new byte[0]));

    final Map<FlowFile, Integer> msgCounts = new HashMap<>();
    msgCounts.put(flowFiles.get(0), 0);

    final PublishResult result = createPublishResult(msgCounts, new HashSet<>(flowFiles), Collections.emptyMap());

    when(mockLease.complete()).thenReturn(result);

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_0_11.REL_SUCCESS, 1);

    verify(mockLease, times(1)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();

    final MockFlowFile mff = runner.getFlowFilesForRelationship(PublishKafkaRecord_0_11.REL_SUCCESS).get(0);
    mff.assertAttributeEquals("msg.count", "0");
}
 

@Test
public void testMultipleSuccess() throws IOException {
    final Set<FlowFile> flowFiles = new HashSet<>();
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));


    when(mockLease.complete()).thenReturn(createAllSuccessPublishResult(flowFiles, 1));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_0_10.REL_SUCCESS, 3);

    verify(mockLease, times(3)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();
}
 

@Test
public void testDelSelectionsJob() {
  JobMeta jobMeta = mock( JobMeta.class );
  Spoon spoon = mock( Spoon.class );
  JobEntryCopy selected1 = mock( JobEntryCopy.class );
  JobEntryCopy selected2 = mock( JobEntryCopy.class );
  when( jobMeta.getSelectedEntries() ).thenReturn( Arrays.asList( selected1, selected2 ) );

  JobGraph jobGraph = mock( JobGraph.class );
  doCallRealMethod().when( jobGraph ).setJobMeta( any( JobMeta.class ) );
  doCallRealMethod().when( jobGraph ).setSpoon( any( Spoon.class ) );
  doCallRealMethod().when( jobGraph ).delSelected( any( JobEntryCopy.class ) );
  jobGraph.setJobMeta( jobMeta );
  jobGraph.setSpoon( spoon );

  jobGraph.delSelected( null );
  verify( spoon ).deleteJobEntryCopies( eq( jobMeta ),
      AdditionalMatchers.aryEq( new JobEntryCopy[] { selected1, selected2 } ) );
}
 

@Test
public void testWriteKeyStoreFailure() throws IOException, GeneralSecurityException {
    setUnlimitedCrypto(false);
    String testPassword = "testPassword";
    IOException ioException = new IOException("Fail");
    doThrow(ioException).when(keyStoreSpi).engineStore(eq(tmpFileOutputStream), AdditionalMatchers.aryEq(testPassword.toCharArray()));
    try {
        TlsHelper.writeKeyStore(keyStore, outputStreamFactory, file, testPassword, true);
        fail("Expected " + ioException);
    } catch (IOException e) {
        assertEquals(ioException, e);
    }
}
 

@Test
public void should_scream_when_no_matchers_inside_not() {
    try {
        mock.simpleMethod(AdditionalMatchers.not("jkl"));
        fail();
    } catch (InvalidUseOfMatchersException e) {
        assertThat(e.getMessage())
                .contains("No matchers found for")
                .containsIgnoringCase("Not(?)");
    }
}
 

@Test
public void testWriteKeyStoreNoTruncate() throws IOException, GeneralSecurityException {
    setUnlimitedCrypto(false);
    String testPassword = "testPassword";
    IOException ioException = new IOException(TlsHelper.ILLEGAL_KEY_SIZE);
    doThrow(ioException).when(keyStoreSpi).engineStore(eq(tmpFileOutputStream), AdditionalMatchers.aryEq(testPassword.toCharArray()));
    try {
        TlsHelper.writeKeyStore(keyStore, outputStreamFactory, file, testPassword, false);
        fail("Expected " + GeneralSecurityException.class);
    } catch (GeneralSecurityException e) {
        assertTrue("Expected exception to contain " + TlsHelper.JCE_URL, e.getMessage().contains(TlsHelper.JCE_URL));
    }
}
 

@Test
public void testSingleSuccess() throws IOException {
    final MockFlowFile flowFile = runner.enqueue("John Doe, 48");

    when(mockLease.complete()).thenReturn(createAllSuccessPublishResult(flowFile, 1));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_2_0.REL_SUCCESS, 1);

    verify(mockLease, times(1)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();
}
 

@Test
public void testSingleFailure() throws IOException {
    final MockFlowFile flowFile = runner.enqueue("John Doe, 48");

    when(mockLease.complete()).thenReturn(createFailurePublishResult(flowFile));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_2_0.REL_FAILURE, 1);

    verify(mockLease, times(1)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(1)).close();
}
 

@Test
public void testSearchUsesRootGroupAsActiveIfNotProvided() {
    // given
    final ControllerFacade testSubject = givenTestSubject();

    // when
    testSubject.search(SEARCH_LITERAL, null);

    // then
    Mockito.verify(searchQueryParser, Mockito.times(1))
            .parse(Mockito.eq(SEARCH_LITERAL), AdditionalMatchers.or(Mockito.any(NiFiUser.class), Mockito.isNull()), Mockito.same(rootGroup), Mockito.same(rootGroup));
}
 

@Test
public void testSliderDrag_ListenerShouldBeCalled() {
  // Lays out slider.
  helper.addContentView(activity);

  // Drag starts from one quarter to the left end to the middle.
  dragSliderBetweenValues(
      slider, SLIDER_VALUE_FROM + SLIDER_VALUE_RANGE / 4, SLIDER_VALUE_TO / 2, 100);

  // Verifies listener calls.
  for (int value = 25; value <= 50; value++) {
    verify(mockOnChangeListener, times(1))
        .onValueChange(eq(slider), AdditionalMatchers.eq((float) value, FLOAT_ERROR), eq(true));
  }
}
 

@Test
public void postTest() throws Exception {
    Mockito.when(req.getParameter("indexingEnabled")).thenReturn("" + GlobalSettings.INDEXING_ENABLED_DEFAULT);
    Mockito.when(req.getParameter("maxConcurrentIndexing")).thenReturn(
        "" + GlobalSettings.MAX_CONCURRENT_INDEXING_DEFAULT);
    Mockito.when(req.getParameter("maxFileSize")).thenReturn("" + GlobalSettings.MAX_FILE_SIZE_DEFAULT);
    Mockito.when(req.getParameter("searchTimeout")).thenReturn("" + GlobalSettings.SEARCH_TIMEOUT_DEFAULT);
    Mockito.when(req.getParameter("noHighlightExtensions")).thenReturn(
        "" + GlobalSettings.NO_HIGHLIGHT_EXTENSIONS_DEFAULT);
    Mockito.when(req.getParameter("maxPreviewLines")).thenReturn("" + GlobalSettings.MAX_PREVIEW_LINES_DEFAULT);
    Mockito.when(req.getParameter("maxMatchLines")).thenReturn("" + GlobalSettings.MAX_MATCH_LINES_DEFAULT);
    Mockito.when(req.getParameter("maxFragments")).thenReturn("" + GlobalSettings.MAX_FRAGMENTS_DEFAULT);
    Mockito.when(req.getParameter("pageSize")).thenReturn("" + GlobalSettings.PAGE_SIZE_DEFAULT);
    Mockito.when(req.getParameter("commitHashBoost")).thenReturn("" + GlobalSettings.COMMIT_HASH_BOOST_DEFAULT);
    Mockito.when(req.getParameter("commitSubjectBoost")).thenReturn(
        "" + GlobalSettings.COMMIT_SUBJECT_BOOST_DEFAULT);
    Mockito.when(req.getParameter("commitBodyBoost")).thenReturn("" + GlobalSettings.COMMIT_BODY_BOOST_DEFAULT);
    Mockito.when(req.getParameter("fileNameBoost")).thenReturn("" + GlobalSettings.FILE_NAME_BOOST_DEFAULT);

    servlet.doPost(req, res);

    Mockito.verify(sm).setGlobalSettings(
        Mockito.eq(GlobalSettings.INDEXING_ENABLED_DEFAULT),
        Mockito.eq(GlobalSettings.MAX_CONCURRENT_INDEXING_DEFAULT),
        Mockito.eq(GlobalSettings.MAX_FILE_SIZE_DEFAULT),
        Mockito.eq(GlobalSettings.SEARCH_TIMEOUT_DEFAULT),
        Mockito.eq(GlobalSettings.NO_HIGHLIGHT_EXTENSIONS_DEFAULT),
        Mockito.eq(GlobalSettings.MAX_PREVIEW_LINES_DEFAULT),
        Mockito.eq(GlobalSettings.MAX_MATCH_LINES_DEFAULT),
        Mockito.eq(GlobalSettings.MAX_FRAGMENTS_DEFAULT),
        Mockito.eq(GlobalSettings.PAGE_SIZE_DEFAULT),
        AdditionalMatchers.eq(GlobalSettings.COMMIT_HASH_BOOST_DEFAULT, 1E-9),
        AdditionalMatchers.eq(GlobalSettings.COMMIT_SUBJECT_BOOST_DEFAULT, 1E-9),
        AdditionalMatchers.eq(GlobalSettings.COMMIT_BODY_BOOST_DEFAULT, 1E-9),
        AdditionalMatchers.eq(GlobalSettings.FILE_NAME_BOOST_DEFAULT, 1E-9));
    Mockito.verify(res).setContentType(Mockito.contains("text/html"));
}
 

@Test
public void testSingleFailure() throws IOException {
    final MockFlowFile flowFile = runner.enqueue("John Doe, 48");

    when(mockLease.complete()).thenReturn(createFailurePublishResult(flowFile));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_0_10.REL_FAILURE, 1);

    verify(mockLease, times(1)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(1)).close();
}
 

private ReceiverAdminClient mockReceiverClientFailOnStopAndSync() throws IOException {
    ReceiverAdminClient receiverAdminClient = mock(ReceiverAdminClient.class);

    ConsumerStatsResponse consumerStatsResponse = new ConsumerStatsResponse();
    consumerStatsResponse.setConsumePosition(positionRs1);
    consumerStatsResponse.setCubeName(cubeName);
    when(receiverAdminClient.pauseConsumers(eq(n4), any(PauseConsumersRequest.class)))
            .thenThrow(new IOException("Mock Receiver Error"));
    when(receiverAdminClient.pauseConsumers(AdditionalMatchers.not(Matchers.eq(n4)),
            any(PauseConsumersRequest.class))).thenReturn(consumerStatsResponse);

    Mockito.doThrow(new IOException()).when(receiverAdminClient).startConsumers(eq(n1),
            any(StartConsumersRequest.class));
    return receiverAdminClient;
}
 

@Test
public void testMultipleMessagesPerFlowFile() throws IOException {
    final List<FlowFile> flowFiles = new ArrayList<>();
    flowFiles.add(runner.enqueue("John Doe, 48\nJane Doe, 47"));
    flowFiles.add(runner.enqueue("John Doe, 48\nJane Doe, 29"));

    final Map<FlowFile, Integer> msgCounts = new HashMap<>();
    msgCounts.put(flowFiles.get(0), 10);
    msgCounts.put(flowFiles.get(1), 20);

    final PublishResult result = createPublishResult(msgCounts, new HashSet<>(flowFiles), Collections.emptyMap());

    when(mockLease.complete()).thenReturn(result);

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_2_0.REL_SUCCESS, 2);

    verify(mockLease, times(2)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(0)).publish(
        any(FlowFile.class), any(Map.class), eq(null), any(byte[].class), eq(TOPIC_NAME), any(InFlightMessageTracker.class), any(Integer.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();

    runner.assertAllFlowFilesContainAttribute("msg.count");
    assertEquals(1, runner.getFlowFilesForRelationship(PublishKafkaRecord_2_0.REL_SUCCESS).stream()
        .filter(ff -> ff.getAttribute("msg.count").equals("10"))
        .count());
    assertEquals(1, runner.getFlowFilesForRelationship(PublishKafkaRecord_2_0.REL_SUCCESS).stream()
        .filter(ff -> ff.getAttribute("msg.count").equals("20"))
        .count());
}
 

@SuppressWarnings("unchecked")
@Test
public void testSingleFailure() throws IOException {
    final MockFlowFile flowFile = runner.enqueue("John Doe, 48");

    when(mockLease.complete()).thenReturn(createFailurePublishResult(flowFile));

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_1_0.REL_FAILURE, 1);

    verify(mockLease, times(1)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(1)).close();
}
 

@Test
public void testIOExceptionWhileWritingShouldRouteToRetry() throws InitializationException, IOException, SchemaNotFoundException {
    configure(proc);

    final RecordSetWriter recordSetWriter = Mockito.mock(RecordSetWriter.class);
    when(recordSetWriter.write(any(Record.class))).thenThrow(new IOException("IOException"));

    final RecordSetWriterFactory recordSetWriterFactory = Mockito.mock(RecordSetWriterFactory.class);
    when(recordSetWriterFactory.getIdentifier()).thenReturn("mock-writer-factory");
    when(recordSetWriterFactory.createWriter(any(ComponentLog.class), AdditionalMatchers.or(any(RecordSchema.class), isNull()), any(OutputStream.class), any(FlowFile.class)))
            .thenReturn(recordSetWriter);

    testRunner.addControllerService("mock-writer-factory", recordSetWriterFactory);
    testRunner.enableControllerService(recordSetWriterFactory);
    testRunner.setProperty(FetchParquet.RECORD_WRITER, "mock-writer-factory");

    final File parquetDir = new File(DIRECTORY);
    final File parquetFile = new File(parquetDir,"testFetchParquetToCSV.parquet");
    final int numUsers = 10;
    writeParquetUsers(parquetFile, numUsers);

    final Map<String,String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.PATH.key(), parquetDir.getAbsolutePath());
    attributes.put(CoreAttributes.FILENAME.key(), parquetFile.getName());

    testRunner.enqueue("TRIGGER", attributes);
    testRunner.run();
    testRunner.assertAllFlowFilesTransferred(FetchParquet.REL_RETRY, 1);

    final MockFlowFile flowFile = testRunner.getFlowFilesForRelationship(FetchParquet.REL_RETRY).get(0);
    flowFile.assertContentEquals("TRIGGER");
}
 

@Test
public void testMultipleMessagesPerFlowFile() throws IOException {
    final List<FlowFile> flowFiles = new ArrayList<>();
    flowFiles.add(runner.enqueue("John Doe, 48\nJane Doe, 47"));
    flowFiles.add(runner.enqueue("John Doe, 48\nJane Doe, 29"));

    final Map<FlowFile, Integer> msgCounts = new HashMap<>();
    msgCounts.put(flowFiles.get(0), 10);
    msgCounts.put(flowFiles.get(1), 20);

    final PublishResult result = createPublishResult(msgCounts, new HashSet<>(flowFiles), Collections.emptyMap());

    when(mockLease.complete()).thenReturn(result);

    runner.run();
    runner.assertAllFlowFilesTransferred(PublishKafkaRecord_0_11.REL_SUCCESS, 2);

    verify(mockLease, times(2)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME));
    verify(mockLease, times(0)).publish(any(FlowFile.class), any(Map.class), eq(null), any(byte[].class), eq(TOPIC_NAME), any(InFlightMessageTracker.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(0)).poison();
    verify(mockLease, times(1)).close();

    runner.assertAllFlowFilesContainAttribute("msg.count");
    assertEquals(1, runner.getFlowFilesForRelationship(PublishKafkaRecord_0_11.REL_SUCCESS).stream()
        .filter(ff -> ff.getAttribute("msg.count").equals("10"))
        .count());
    assertEquals(1, runner.getFlowFilesForRelationship(PublishKafkaRecord_0_11.REL_SUCCESS).stream()
        .filter(ff -> ff.getAttribute("msg.count").equals("20"))
        .count());
}
 

@SuppressWarnings("unchecked")
@Test
public void testSomeSuccessSomeFailure() throws IOException {
    final List<FlowFile> flowFiles = new ArrayList<>();
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));
    flowFiles.add(runner.enqueue("John Doe, 48"));

    final Map<FlowFile, Integer> msgCounts = new HashMap<>();
    msgCounts.put(flowFiles.get(0), 10);
    msgCounts.put(flowFiles.get(1), 20);

    final Map<FlowFile, Exception> failureMap = new HashMap<>();
    failureMap.put(flowFiles.get(2), new RuntimeException("Intentional Unit Test Exception"));
    failureMap.put(flowFiles.get(3), new RuntimeException("Intentional Unit Test Exception"));

    final PublishResult result = createPublishResult(msgCounts, new HashSet<>(flowFiles.subList(0, 2)), failureMap);

    when(mockLease.complete()).thenReturn(result);

    runner.run();
    runner.assertTransferCount(PublishKafkaRecord_1_0.REL_SUCCESS, 0);
    runner.assertTransferCount(PublishKafkaRecord_1_0.REL_FAILURE, 4);

    verify(mockLease, times(4)).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class),
            AdditionalMatchers.or(any(RecordSchema.class), isNull()), eq(null), eq(TOPIC_NAME), nullable(Function.class));
    verify(mockLease, times(1)).complete();
    verify(mockLease, times(1)).close();

    assertTrue(runner.getFlowFilesForRelationship(PublishKafkaRecord_1_0.REL_FAILURE).stream()
        .noneMatch(ff -> ff.getAttribute("msg.count") != null));
}
 

@Test
public void testWriteKeyStoreFailure() throws IOException, GeneralSecurityException {
    setUnlimitedCrypto(false);
    String testPassword = "testPassword";
    IOException ioException = new IOException("Fail");
    doThrow(ioException).when(keyStoreSpi).engineStore(eq(tmpFileOutputStream), AdditionalMatchers.aryEq(testPassword.toCharArray()));
    try {
        TlsHelper.writeKeyStore(keyStore, outputStreamFactory, file, testPassword, true);
        fail("Expected " + ioException);
    } catch (IOException e) {
        assertEquals(ioException, e);
    }
}