org.apache.nifi.provenance.ProvenanceEventRecord Java Examples

private void rebuildIndex(final File tempIndexDir, final File migratedIndexDir) throws IOException {
    final EventIndexWriter writer = indexManager.borrowIndexWriter(tempIndexDir);

    try {
        final EventIterator eventIterator = eventStore.getEventsByTimestamp(minTimestamp, maxTimestamp);

        final StopWatch stopWatch = new StopWatch(true);

        Optional<ProvenanceEventRecord> optionalEvent;
        while ((optionalEvent = eventIterator.nextEvent()).isPresent()) {
            final ProvenanceEventRecord event = optionalEvent.get();

            final Document document = eventConverter.convert(event, event.getEventId());
            writer.index(document, Integer.MAX_VALUE);
            successCount++;
        }

        writer.commit();
        stopWatch.stop();
        logger.info("Successfully indexed {} events to {} in {}", successCount, tempIndexDir, stopWatch.getDuration());
    } finally {
        indexManager.returnIndexWriter(writer, true, true);
    }

    Files.move(tempIndexDir.toPath(), migratedIndexDir.toPath());
}
 

@Override
public void associate(final FlowFile flowFile, final String alternateIdentifierNamespace, final String alternateIdentifier) {
    try {
        String trimmedNamespace = alternateIdentifierNamespace.trim();
        if (trimmedNamespace.endsWith(":")) {
            trimmedNamespace = trimmedNamespace.substring(0, trimmedNamespace.length() - 1);
        }

        String trimmedIdentifier = alternateIdentifier.trim();
        if (trimmedIdentifier.startsWith(":")) {
            if (trimmedIdentifier.length() == 1) {
                throw new IllegalArgumentException("Illegal alternateIdentifier: " + alternateIdentifier);
            }
            trimmedIdentifier = trimmedIdentifier.substring(1);
        }

        final String alternateIdentifierUri = trimmedNamespace + ":" + trimmedIdentifier;
        final ProvenanceEventRecord record = build(flowFile, ProvenanceEventType.ADDINFO).setAlternateIdentifierUri(alternateIdentifierUri).build();
        events.add(record);
    } catch (final Exception e) {
        logger.error("Failed to generate Provenance Event due to " + e);
        if (logger.isDebugEnabled()) {
            logger.error("", e);
        }
    }
}
 

/**
 * Get the provenance event with the specified event id.
 *
 * @param eventId event id
 * @return the provenance event with the specified event id
 */
public ProvenanceEventDTO getProvenanceEvent(final Long eventId) {
    try {
        final ProvenanceEventRecord event = flowController.getProvenanceRepository().getEvent(eventId);
        if (event == null) {
            throw new ResourceNotFoundException("Unable to find the specified event.");
        }

        // get the flowfile attributes and authorize the event
        final Map<String, String> attributes = event.getAttributes();
        final Authorizable dataAuthorizable;
        if (event.isRemotePortType()) {
            dataAuthorizable = flowController.createRemoteDataAuthorizable(event.getComponentId());
        } else {
            dataAuthorizable = flowController.createLocalDataAuthorizable(event.getComponentId());
        }
        dataAuthorizable.authorize(authorizer, RequestAction.READ, NiFiUserUtils.getNiFiUser(), attributes);

        // convert the event
        return createProvenanceEventDto(event, false);
    } catch (final IOException ioe) {
        throw new NiFiCoreException("An error occurred while getting the specified event.", ioe);
    }
}
 

@Override
public Optional<ProvenanceEventRecord> nextEvent() throws IOException {
    while (true) {
        final Optional<ProvenanceEventRecord> next = iterator.nextEvent();
        if (!next.isPresent()) {
            return next;
        }

        if (authorizer.isAuthorized(next.get())) {
            return next;
        }

        final Optional<ProvenanceEventRecord> eventOption = transformer.transform(next.get());
        if (eventOption.isPresent()) {
            return eventOption;
        }
    }
}
 

@Test
public void testCreateEventIfNewFlowFileWithoutReceive() throws ExecutionException, InterruptedException, IOException {
    final ProcessorNode createProcessor = createProcessorNode((context, session) -> {
        FlowFile flowFile = session.create();

        final Map<String, String> attrs = new HashMap<>();
        attrs.put("test", "integration");
        attrs.put("integration", "true");

        flowFile = session.putAllAttributes(flowFile, attrs);
        session.transfer(flowFile, REL_SUCCESS);
    }, REL_SUCCESS);

    connect(createProcessor, getTerminateAllProcessor(), REL_SUCCESS);
    triggerOnce(createProcessor);

    // There should be exactly 1 event.
    final ProvenanceEventRepository provRepo = getProvenanceRepository();
    assertEquals(0L, provRepo.getMaxEventId().longValue());

    final ProvenanceEventRecord firstEvent = provRepo.getEvent(0L);
    assertEquals(ProvenanceEventType.CREATE, firstEvent.getEventType());
    assertEquals("integration", firstEvent.getAttribute("test"));
    assertEquals("true", firstEvent.getAttribute("integration"));
}
 

private ProvenanceEventRecord getRecord(final Document d, final RecordReader reader) throws IOException {
    final IndexableField blockField = d.getField(FieldNames.BLOCK_INDEX);
    if ( blockField == null ) {
        reader.skipTo(getByteOffset(d, reader));
    } else {
        reader.skipToBlock(blockField.numericValue().intValue());
    }

    StandardProvenanceEventRecord record;
    while ( (record = reader.nextRecord()) != null) {
        final IndexableField idField = d.getField(SearchableFields.Identifier.getSearchableFieldName());
        if ( idField == null || idField.numericValue().longValue() == record.getEventId() ) {
            break;
        }
    }

    if (record == null) {
        logger.warn("Failed to read Provenance Event for '" + d + "'. The event file may be missing or corrupted");
    }

    return record;
}
 

@Override
public long getMaxEventId() throws IOException {
    if (tocReader != null) {
        final long lastBlockOffset = tocReader.getLastBlockOffset();
        skipToBlock(tocReader.getBlockIndex(lastBlockOffset));
    }

    ProvenanceEventRecord record;
    ProvenanceEventRecord lastRecord = null;
    try {
        while ((record = nextRecord()) != null) {
            lastRecord = record;
        }
    } catch (final EOFException eof) {
        // This can happen if we stop NIFi while the record is being written.
        // This is OK, we just ignore this record. The session will not have been
        // committed, so we can just process the FlowFile again.
    }

    return lastRecord == null ? -1L : lastRecord.getEventId();
}
 

@Test
public void testAttributesModifiedEmitted() throws IOException {
    final FlowFileRecord flowFile = new StandardFlowFileRecord.Builder()
            .id(1L)
            .addAttribute("uuid", "000000000000-0000-0000-0000-00000000")
            .build();
    this.flowFileQueue.put(flowFile);

    FlowFile existingFlowFile = session.get();
    existingFlowFile = session.putAttribute(existingFlowFile, "attr", "a");
    session.transfer(existingFlowFile, new Relationship.Builder().name("A").build());
    session.commit();

    final List<ProvenanceEventRecord> events = provenanceRepo.getEvents(0L, 10000);
    assertFalse(events.isEmpty());
    assertEquals(1, events.size());

    final ProvenanceEventRecord event = events.get(0);
    assertEquals(ProvenanceEventType.ATTRIBUTES_MODIFIED, event.getEventType());
}
 

ProvenanceEventRecord drop(final FlowFile flowFile, final String reason) {
    try {
        final ProvenanceEventBuilder builder = build(flowFile, ProvenanceEventType.DROP);
        if (reason != null) {
            builder.setDetails("Discard reason: " + reason);
        }
        final ProvenanceEventRecord record = builder.build();
        events.add(record);
        return record;
    } catch (final Exception e) {
        logger.error("Failed to generate Provenance Event due to " + e);
        if (logger.isDebugEnabled()) {
            logger.error("", e);
        }
        return null;
    }
}
 

/**
 * Authorizes access to replay a specified provenance event.
 *
 * @param event event
 */
private AuthorizationResult checkAuthorizationForReplay(final ProvenanceEventRecord event) {
    // if the connection id isn't specified, then the replay wouldn't be available anyways and we have nothing to authorize against so deny it`
    if (event.getSourceQueueIdentifier() == null) {
        return AuthorizationResult.denied("The connection id in the provenance event is unknown.");
    }

    final NiFiUser user = NiFiUserUtils.getNiFiUser();
    final Authorizable dataAuthorizable;
    if (event.isRemotePortType()) {
        dataAuthorizable = flowController.createRemoteDataAuthorizable(event.getComponentId());
    } else {
        dataAuthorizable = flowController.createLocalDataAuthorizable(event.getComponentId());
    }

    final Map<String, String> eventAttributes = event.getAttributes();

    // ensure we can read the data
    final AuthorizationResult result = dataAuthorizable.checkAuthorization(authorizer, RequestAction.READ, user, eventAttributes);
    if (!Result.Approved.equals(result.getResult())) {
        return result;
    }

    // ensure we can write the data
    return dataAuthorizable.checkAuthorization(authorizer, RequestAction.WRITE, user, eventAttributes);
}
 

private ProvenanceEventRecord createProvenanceEventRecord(final String componentId, final String componentType) {
    final String uuid = "10000000-0000-0000-0000-000000000000";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("abc", "xyz");
    attributes.put("xyz", "abc");
    attributes.put("filename", "file-" + uuid);
    attributes.put("emptyVal",null);

    final Map<String, String> prevAttrs = new HashMap<>();
    attributes.put("filename", "1234.xyz");

    final ProvenanceEventBuilder builder = new StandardProvenanceEventRecord.Builder();
    builder.setEventTime(System.currentTimeMillis());
    builder.setEventType(ProvenanceEventType.RECEIVE);
    builder.setTransitUri("nifi://unit-test");
    attributes.put("uuid", uuid);
    builder.fromFlowFile(createFlowFile(3L, attributes));
    builder.setAttributes(prevAttrs, attributes);
    builder.setComponentId(componentId);
    builder.setComponentType(componentType);
    return builder.build();
}
 

@Test
public void testNoTimeLimit() throws InitializationException, ClassNotFoundException, SQLException, IOException {
    invokeOnTrigger(QUERY_WITH_EL, true, "Avro");

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(3, provenanceEvents.size());

    final ProvenanceEventRecord provenance0 = provenanceEvents.get(0);
    assertEquals(ProvenanceEventType.FORK, provenance0.getEventType());

    final ProvenanceEventRecord provenance1 = provenanceEvents.get(1);
    assertEquals(ProvenanceEventType.FETCH, provenance1.getEventType());
    assertEquals("jdbc:derby:target/db;create=true", provenance1.getTransitUri());

    final ProvenanceEventRecord provenance2 = provenanceEvents.get(2);
    assertEquals(ProvenanceEventType.FORK, provenance2.getEventType());
}
 

private ProvenanceEventRecord createEvent(final long timestamp, final String uuid) {
    final Map<String, String> previousAttributes = new HashMap<>();
    previousAttributes.put("uuid", uuid);
    final Map<String, String> updatedAttributes = new HashMap<>();
    updatedAttributes.put("updated", "true");

    final ProvenanceEventRecord event = new StandardProvenanceEventRecord.Builder()
        .setEventType(ProvenanceEventType.CONTENT_MODIFIED)
        .setAttributes(previousAttributes, updatedAttributes)
        .setComponentId("component-1")
        .setComponentType("unit test")
        .setEventId(idGenerator.getAndIncrement())
        .setEventTime(timestamp)
        .setFlowFileEntryDate(timestamp)
        .setFlowFileUUID(uuid)
        .setLineageStartDate(timestamp)
        .setCurrentContentClaim("container", "section", "unit-test-id", 0L, 1024L)
        .build();

    return event;
}
 

@Override
public void associate(final FlowFile flowFile, final String alternateIdentifierNamespace, final String alternateIdentifier) {
    try {
        String trimmedNamespace = alternateIdentifierNamespace.trim();
        if (trimmedNamespace.endsWith(":")) {
            trimmedNamespace = trimmedNamespace.substring(0, trimmedNamespace.length() - 1);
        }

        String trimmedIdentifier = alternateIdentifier.trim();
        if (trimmedIdentifier.startsWith(":")) {
            if (trimmedIdentifier.length() == 1) {
                throw new IllegalArgumentException("Illegal alternateIdentifier: " + alternateIdentifier);
            }
            trimmedIdentifier = trimmedIdentifier.substring(1);
        }

        final String alternateIdentifierUri = trimmedNamespace + ":" + trimmedIdentifier;
        final ProvenanceEventRecord record = build(flowFile, ProvenanceEventType.ADDINFO).setAlternateIdentifierUri(alternateIdentifierUri).build();
        events.add(record);
    } catch (final Exception e) {
        logger.error("Failed to generate Provenance Event due to " + e);
        if (logger.isDebugEnabled()) {
            logger.error("", e);
        }
    }
}
 

private long updateLastEventId(final List<ProvenanceEventRecord> events, final StateManager stateManager) {
    if (events == null || events.isEmpty()) {
        return firstEventId;
    }

    // Store the id of the last event so we know where we left off
    final ProvenanceEventRecord lastEvent = events.get(events.size() - 1);
    final String lastEventId = String.valueOf(lastEvent.getEventId());
    try {
        Map<String, String> newMapOfState = new HashMap<>();
        newMapOfState.put(LAST_EVENT_ID_KEY, lastEventId);
        stateManager.setState(newMapOfState, Scope.LOCAL);
    } catch (final IOException ioe) {
        logger.error("Failed to update state to {} due to {}; this could result in events being re-sent after a restart. The message of {} was: {}",
                new Object[]{lastEventId, ioe, ioe, ioe.getMessage()}, ioe);
    }

    return lastEvent.getEventId() + 1;
}
 

/**
 * Tests basic ES functionality against a local or test ES cluster
 */
@Test
@Ignore("Comment this out if you want to run against local or test ES")
public void testPutElasticSearchBasic() {
    System.out.println("Starting test " + new Object() {
    }.getClass().getEnclosingMethod().getName());
    final TestRunner runner = TestRunners.newTestRunner(new PutElasticsearchHttpRecord());

    runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200");
    runner.setProperty(PutElasticsearchHttpRecord.INDEX, "doc");
    runner.setProperty(PutElasticsearchHttpRecord.TYPE, "status");
    runner.setProperty(PutElasticsearchHttpRecord.ID_RECORD_PATH, "/id");
    runner.assertValid();

    runner.enqueue(new byte[0], new HashMap<String, String>() {{
        put("doc_id", "28039652140");
    }});

    runner.enqueue(new byte[0]);
    runner.run(1, true, true);
    runner.assertAllFlowFilesTransferred(PutElasticsearchHttpRecord.REL_SUCCESS, 1);
    List<ProvenanceEventRecord> provEvents = runner.getProvenanceEvents();
    assertNotNull(provEvents);
    assertEquals(1, provEvents.size());
    assertEquals(ProvenanceEventType.SEND, provEvents.get(0).getEventType());
}
 

@Test
public void testNoTimeLimit() throws InitializationException, ClassNotFoundException, SQLException, IOException {
    invokeOnTrigger(QUERY_WITH_EL, true, "Avro");

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(3, provenanceEvents.size());

    final ProvenanceEventRecord provenance0 = provenanceEvents.get(0);
    assertEquals(ProvenanceEventType.FORK, provenance0.getEventType());

    final ProvenanceEventRecord provenance1 = provenanceEvents.get(1);
    assertEquals(ProvenanceEventType.FETCH, provenance1.getEventType());
    assertEquals("jdbc:derby:target/db;create=true", provenance1.getTransitUri());

    final ProvenanceEventRecord provenance2 = provenanceEvents.get(2);
    assertEquals(ProvenanceEventType.FORK, provenance2.getEventType());
}
 

@Test
public void testFilterComponentTypeSuccess() throws IOException, InitializationException {
    final Map<PropertyDescriptor, String> properties = new HashMap<>();
    for (final PropertyDescriptor descriptor : new MockSiteToSiteProvenanceReportingTask().getSupportedPropertyDescriptors()) {
        properties.put(descriptor, descriptor.getDefaultValue());
    }
    properties.put(SiteToSiteUtils.BATCH_SIZE, "1000");
    properties.put(SiteToSiteProvenanceReportingTask.FILTER_COMPONENT_TYPE, "dummy.*");

    ProvenanceEventRecord event = createProvenanceEventRecord();

    MockSiteToSiteProvenanceReportingTask task = setup(event, properties);
    task.initialize(initContext);
    task.onScheduled(confContext);
    task.onTrigger(context);

    assertEquals(3, task.dataSent.size());
}
 

@Test
public void testFilterComponentIdSuccess() throws IOException, InitializationException {
    final Map<PropertyDescriptor, String> properties = new HashMap<>();
    for (final PropertyDescriptor descriptor : new MockSiteToSiteProvenanceReportingTask().getSupportedPropertyDescriptors()) {
        properties.put(descriptor, descriptor.getDefaultValue());
    }
    properties.put(SiteToSiteUtils.BATCH_SIZE, "1000");
    properties.put(SiteToSiteProvenanceReportingTask.FILTER_COMPONENT_ID, "2345, 5678,  1234");

    ProvenanceEventRecord event = createProvenanceEventRecord();

    MockSiteToSiteProvenanceReportingTask task = setup(event, properties);
    task.initialize(initContext);
    task.onScheduled(confContext);
    task.onTrigger(context);

    assertEquals(3, task.dataSent.size());
}
 

private void analyzeLineagePath(AnalysisContext analysisContext, LineagePath lineagePath) {
    final List<ProvenanceEventRecord> events = lineagePath.getEvents();

    final DataSetRefs parentRefs = new DataSetRefs(events.get(0).getComponentId());
    events.forEach(event -> {
        final DataSetRefs refs = executeAnalyzer(analysisContext, event);
        if (refs == null || refs.isEmpty()) {
            return;
        }
        refs.getInputs().forEach(parentRefs::addInput);
        refs.getOutputs().forEach(parentRefs::addOutput);
    });

    lineagePath.setRefs(parentRefs);

    // Analyse parents.
    lineagePath.getParents().forEach(parent -> analyzeLineagePath(analysisContext, parent));
}
 

private ProvenanceEventRecord createEvent() {
    final String uuid = UUID.randomUUID().toString();
    final Map<String, String> previousAttributes = new HashMap<>();
    previousAttributes.put("uuid", uuid);
    final Map<String, String> updatedAttributes = new HashMap<>();
    updatedAttributes.put("updated", "true");

    return new StandardProvenanceEventRecord.Builder()
        .setEventType(ProvenanceEventType.CONTENT_MODIFIED)
        .setAttributes(previousAttributes, updatedAttributes)
        .setComponentId("component-1")
        .setComponentType("unit test")
        .setEventTime(System.currentTimeMillis())
        .setFlowFileEntryDate(System.currentTimeMillis())
        .setFlowFileUUID(uuid)
        .setLineageStartDate(System.currentTimeMillis())
        .setCurrentContentClaim("container", "section", "unit-test-id", 0L, 1024L)
        .build();
}
 

@Override
public List<ProvenanceEventRecord> getEvents(long firstRecordId, int maxResults, EventAuthorizer authorizer, EventTransformer transformer) throws IOException {
    final List<ProvenanceEventRecord> events = new ArrayList<>();
    for (int i = 0; i < maxResults; i++) {
        final Optional<ProvenanceEventRecord> eventOption = getEvent(firstRecordId + i);
        if (!eventOption.isPresent()) {
            break;
        }

        events.add(eventOption.get());
    }

    return events;
}
 

@Test
public void testSingJsonDocAndExtractedRowId() throws IOException, InitializationException {
    final TestRunner runner = getTestRunner(DEFAULT_TABLE_NAME, DEFAULT_COLUMN_FAMILY, "1");
    final MockHBaseClientService hBaseClient = getHBaseClientService(runner);
    runner.setProperty(PutHBaseJSON.ROW_FIELD_NAME, "rowField");

    final String content = "{ \"rowField\" : \"myRowId\", \"field1\" : \"value1\", \"field2\" : \"value2\" }";
    runner.enqueue(content.getBytes(StandardCharsets.UTF_8));
    runner.run();
    runner.assertAllFlowFilesTransferred(PutHBaseCell.REL_SUCCESS);

    final MockFlowFile outFile = runner.getFlowFilesForRelationship(PutHBaseCell.REL_SUCCESS).get(0);
    outFile.assertContentEquals(content);

    assertNotNull(hBaseClient.getFlowFilePuts());
    assertEquals(1, hBaseClient.getFlowFilePuts().size());

    final List<PutFlowFile> puts = hBaseClient.getFlowFilePuts().get(DEFAULT_TABLE_NAME);
    assertEquals(1, puts.size());

    // should be a put with row id of myRowId, and rowField shouldn't end up in the columns
    final Map<String,byte[]> expectedColumns1 = new HashMap<>();
    expectedColumns1.put("field1", hBaseClient.toBytes("value1"));
    expectedColumns1.put("field2", hBaseClient.toBytes("value2"));
    HBaseTestUtil.verifyPut("myRowId", DEFAULT_COLUMN_FAMILY, expectedColumns1, puts);

    final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
    assertEquals(1, events.size());
    HBaseTestUtil.verifyEvent(runner.getProvenanceEvents(), "hbase://" + DEFAULT_TABLE_NAME + "/myRowId", ProvenanceEventType.SEND);
}
 

private void assertProvenanceEvents(int count){
    List<ProvenanceEventRecord> provenanceEvents = testRunner.getProvenanceEvents();
    assertNotNull(provenanceEvents);
    assertEquals(count, provenanceEvents.size());
    if (count > 0) {
        assertEquals(ProvenanceEventType.RECEIVE, provenanceEvents.get(0).getEventType());
    }
}
 

@Test
public void testValidMessageStaticPropertiesTcp() {
    final String pri = "34";
    final String version = "1";
    final String stamp = "2003-10-11T22:14:15.003Z";
    final String host = "mymachine.example.com";
    final String body = "su - ID47 - BOM'su root' failed for lonvick on /dev/pts/8";

    final String expectedMessage = "<" + pri + ">" + version + " " + stamp + " " + host + " " + body;

    runner.setProperty(PutSyslog.PROTOCOL, PutSyslog.TCP_VALUE);
    runner.setProperty(PutSyslog.MSG_PRIORITY, pri);
    runner.setProperty(PutSyslog.MSG_VERSION, version);
    runner.setProperty(PutSyslog.MSG_TIMESTAMP, stamp);
    runner.setProperty(PutSyslog.MSG_HOSTNAME, host);
    runner.setProperty(PutSyslog.MSG_BODY, body);

    runner.enqueue("incoming data".getBytes(Charset.forName("UTF-8")));
    runner.run();

    runner.assertAllFlowFilesTransferred(PutSyslog.REL_SUCCESS, 1);
    Assert.assertEquals(1, sender.messages.size());
    Assert.assertEquals(expectedMessage, sender.messages.get(0).replace("\n", ""));

    final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
    Assert.assertNotNull(events);
    Assert.assertEquals(1, events.size());

    final ProvenanceEventRecord event = events.get(0);
    Assert.assertEquals(ProvenanceEventType.SEND, event.getEventType());
    Assert.assertEquals("TCP://localhost:12345", event.getTransitUri());
}
 

@Override
public ProvenanceEventRecord getProvenanceEvent(long eventId) {
    try {
        return provenanceRepository.getEvent(eventId);
    } catch (IOException e) {
        logger.error("Failed to get provenance event for {} due to {}", new Object[]{eventId, e}, e);
        return null;
    }
}
 

@Override
public List<ProvenanceEventRecord> getEvents(final List<Long> eventIds, final EventAuthorizer authorizer, final EventTransformer transformer) throws IOException {
    if (eventIds == null || eventIds.isEmpty()) {
        return Collections.emptyList();
    }

    return getEvents(eventIds.size(), authorizer, part -> part.createEventIterator(eventIds), transformer);
}
 

@Test
public void testAttributesModifiedNotCreatedIfProcessorEmitsOtherEvent() throws ExecutionException, InterruptedException, IOException {
    final ProcessorNode createProcessor = createGenerateProcessor(0);

    final ProcessorNode updateProcessor = createProcessorNode((context, session) -> {
        FlowFile flowFile = session.get();

        final Map<String, String> attrs = new HashMap<>();
        attrs.put("test", "integration");
        attrs.put("integration", "true");

        flowFile = session.putAllAttributes(flowFile, attrs);
        session.getProvenanceReporter().fetch(flowFile, "nifi://unit.test");
        session.transfer(flowFile, REL_SUCCESS);
    }, REL_SUCCESS);

    connect(createProcessor, updateProcessor, REL_SUCCESS);
    connect(updateProcessor, getTerminateAllProcessor(), REL_SUCCESS);

    triggerOnce(createProcessor);
    triggerOnce(updateProcessor);

    // There should be exactly 1 event.
    final ProvenanceEventRepository provRepo = getProvenanceRepository();
    assertEquals(1L, provRepo.getMaxEventId().longValue());

    final ProvenanceEventRecord firstEvent = provRepo.getEvent(0L);
    assertEquals(ProvenanceEventType.CREATE, firstEvent.getEventType());
    assertNull(firstEvent.getAttribute("test"));
    assertNull(firstEvent.getAttribute("integration"));

    final ProvenanceEventRecord secondEvent = provRepo.getEvent(1L);
    assertEquals(ProvenanceEventType.FETCH, secondEvent.getEventType());
    assertEquals("integration", secondEvent.getAttribute("test"));
    assertEquals("true", secondEvent.getAttribute("integration"));
    assertEquals("nifi://unit.test", secondEvent.getTransitUri());
}
 

private void verifyProvenance(int expectedNumEvents) {
    List<ProvenanceEventRecord> provEvents = runner.getProvenanceEvents();
    Assert.assertEquals(expectedNumEvents, provEvents.size());

    for (ProvenanceEventRecord event : provEvents) {
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue(event.getTransitUri().startsWith("udp://"));
    }
}
 

@Test
public void testGenerateFlowFile() {
    final String processorName = "GenerateFlowFile";
    final String processorId = "processor-1234";
    final ProvenanceEventRecord record = Mockito.mock(ProvenanceEventRecord.class);
    when(record.getComponentType()).thenReturn(processorName);
    when(record.getComponentId()).thenReturn(processorId);
    when(record.getEventType()).thenReturn(ProvenanceEventType.CREATE);

    final NamespaceResolvers namespaceResolvers = Mockito.mock(NamespaceResolvers.class);
    when(namespaceResolvers.fromHostNames(matches(".+\\.example\\.com"))).thenReturn("namespace1");

    final List<ConnectionStatus> connections = new ArrayList<>();

    final AnalysisContext context = Mockito.mock(AnalysisContext.class);
    when(context.getNamespaceResolver()).thenReturn(namespaceResolvers);
    when(context.findConnectionTo(processorId)).thenReturn(connections);
    when(context.getNiFiNamespace()).thenReturn("test_namespace");

    final NiFiProvenanceEventAnalyzer analyzer = NiFiProvenanceEventAnalyzerFactory.getAnalyzer(processorName, null, record.getEventType());
    assertNotNull(analyzer);

    final DataSetRefs refs = analyzer.analyze(context, record);
    assertEquals(1, refs.getInputs().size());
    assertEquals(0, refs.getOutputs().size());
    Referenceable ref = refs.getInputs().iterator().next();
    assertEquals("nifi_data", ref.getTypeName());
    assertEquals("GenerateFlowFile", ref.get(ATTR_NAME));
    assertEquals("processor-1234@test_namespace", ref.get(ATTR_QUALIFIED_NAME));
}