Examples with Connection org.apache.nifi.connectable.Connection used on opensource projects

Example 76 with Connection

use of org.apache.nifi.connectable.Connection in project nifi by apache.

the class TestWriteAheadFlowFileRepository method testRestartWithOneRecord.

@Test
public void testRestartWithOneRecord() throws IOException {
    final Path path = Paths.get("target/test-repo");
    if (Files.exists(path)) {
        FileUtils.deleteFile(path.toFile(), true);
    }
    final WriteAheadFlowFileRepository repo = new WriteAheadFlowFileRepository(NiFiProperties.createBasicNiFiProperties(null, null));
    repo.initialize(new StandardResourceClaimManager());
    final TestQueueProvider queueProvider = new TestQueueProvider();
    repo.loadFlowFiles(queueProvider, 0L);
    final List<FlowFileRecord> flowFileCollection = new ArrayList<>();
    final Connection connection = Mockito.mock(Connection.class);
    when(connection.getIdentifier()).thenReturn("1234");
    final FlowFileQueue queue = Mockito.mock(FlowFileQueue.class);
    when(queue.getIdentifier()).thenReturn("1234");
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(final InvocationOnMock invocation) throws Throwable {
            flowFileCollection.add((FlowFileRecord) invocation.getArguments()[0]);
            return null;
        }
    }).when(queue).put(any(FlowFileRecord.class));
    when(connection.getFlowFileQueue()).thenReturn(queue);
    queueProvider.addConnection(connection);
    StandardFlowFileRecord.Builder ffBuilder = new StandardFlowFileRecord.Builder();
    ffBuilder.id(1L);
    ffBuilder.addAttribute("abc", "xyz");
    ffBuilder.size(0L);
    final FlowFileRecord flowFileRecord = ffBuilder.build();
    final List<RepositoryRecord> records = new ArrayList<>();
    final StandardRepositoryRecord record = new StandardRepositoryRecord(null);
    record.setWorking(flowFileRecord);
    record.setDestination(connection.getFlowFileQueue());
    records.add(record);
    repo.updateRepository(records);
    // update to add new attribute
    ffBuilder = new StandardFlowFileRecord.Builder().fromFlowFile(flowFileRecord).addAttribute("hello", "world");
    final FlowFileRecord flowFileRecord2 = ffBuilder.build();
    record.setWorking(flowFileRecord2);
    repo.updateRepository(records);
    // update size but no attribute
    ffBuilder = new StandardFlowFileRecord.Builder().fromFlowFile(flowFileRecord2).size(40L);
    final FlowFileRecord flowFileRecord3 = ffBuilder.build();
    record.setWorking(flowFileRecord3);
    repo.updateRepository(records);
    repo.close();
    // restore
    final WriteAheadFlowFileRepository repo2 = new WriteAheadFlowFileRepository(NiFiProperties.createBasicNiFiProperties(null, null));
    repo2.initialize(new StandardResourceClaimManager());
    repo2.loadFlowFiles(queueProvider, 0L);
    assertEquals(1, flowFileCollection.size());
    final FlowFileRecord flowFile = flowFileCollection.get(0);
    assertEquals(1L, flowFile.getId());
    assertEquals("xyz", flowFile.getAttribute("abc"));
    assertEquals(40L, flowFile.getSize());
    assertEquals("world", flowFile.getAttribute("hello"));
    repo2.close();
}

Example 77 with Connection

use of org.apache.nifi.connectable.Connection in project nifi by apache.

the class TestConnectableTask method testIsWorkToDo.

@Test
public void testIsWorkToDo() {
    final ProcessorNode procNode = Mockito.mock(ProcessorNode.class);
    Mockito.when(procNode.hasIncomingConnection()).thenReturn(false);
    final Processor processor = Mockito.mock(Processor.class);
    Mockito.when(procNode.getIdentifier()).thenReturn("123");
    Mockito.when(procNode.getRunnableComponent()).thenReturn(processor);
    final FlowController flowController = Mockito.mock(FlowController.class);
    Mockito.when(flowController.getStateManagerProvider()).thenReturn(Mockito.mock(StateManagerProvider.class));
    final RepositoryContext repoContext = Mockito.mock(RepositoryContext.class);
    Mockito.when(repoContext.getFlowFileEventRepository()).thenReturn(Mockito.mock(FlowFileEventRepository.class));
    final RepositoryContextFactory contextFactory = Mockito.mock(RepositoryContextFactory.class);
    Mockito.when(contextFactory.newProcessContext(Mockito.any(Connectable.class), Mockito.any(AtomicLong.class))).thenReturn(repoContext);
    final LifecycleState scheduleState = new LifecycleState();
    final StringEncryptor encryptor = Mockito.mock(StringEncryptor.class);
    ConnectableTask task = new ConnectableTask(Mockito.mock(SchedulingAgent.class), procNode, flowController, contextFactory, scheduleState, encryptor);
    // There is work to do because there are no incoming connections.
    assertFalse(task.invoke().isYield());
    // Test with only a single connection that is self-looping and empty
    final Connection selfLoopingConnection = Mockito.mock(Connection.class);
    when(selfLoopingConnection.getSource()).thenReturn(procNode);
    when(selfLoopingConnection.getDestination()).thenReturn(procNode);
    when(procNode.hasIncomingConnection()).thenReturn(true);
    when(procNode.getIncomingConnections()).thenReturn(Collections.singletonList(selfLoopingConnection));
    assertFalse(task.invoke().isYield());
    // Test with only a single connection that is self-looping and empty
    final FlowFileQueue flowFileQueue = Mockito.mock(FlowFileQueue.class);
    when(flowFileQueue.isActiveQueueEmpty()).thenReturn(true);
    final FlowFileQueue nonEmptyQueue = Mockito.mock(FlowFileQueue.class);
    when(nonEmptyQueue.isActiveQueueEmpty()).thenReturn(false);
    when(selfLoopingConnection.getFlowFileQueue()).thenReturn(nonEmptyQueue);
    assertFalse(task.invoke().isYield());
    // Test with only a non-looping Connection that has no FlowFiles
    final Connection emptyConnection = Mockito.mock(Connection.class);
    when(emptyConnection.getSource()).thenReturn(Mockito.mock(ProcessorNode.class));
    when(emptyConnection.getDestination()).thenReturn(procNode);
    when(emptyConnection.getFlowFileQueue()).thenReturn(flowFileQueue);
    when(procNode.getIncomingConnections()).thenReturn(Collections.singletonList(emptyConnection));
    // Create a new ConnectableTask because we want to have a different value for the 'hasNonLoopConnection' value, which is calculated in the task's constructor.
    task = new ConnectableTask(Mockito.mock(SchedulingAgent.class), procNode, flowController, contextFactory, scheduleState, encryptor);
    assertTrue(task.invoke().isYield());
    // test when the queue has data
    final Connection nonEmptyConnection = Mockito.mock(Connection.class);
    when(nonEmptyConnection.getSource()).thenReturn(Mockito.mock(ProcessorNode.class));
    when(nonEmptyConnection.getDestination()).thenReturn(procNode);
    when(nonEmptyConnection.getFlowFileQueue()).thenReturn(nonEmptyQueue);
    when(procNode.getIncomingConnections()).thenReturn(Collections.singletonList(nonEmptyConnection));
    assertFalse(task.invoke().isYield());
}

Example 78 with Connection

use of org.apache.nifi.connectable.Connection in project nifi by apache.

the class StandardProcessorNode method getUndefinedRelationships.

public Set<Relationship> getUndefinedRelationships() {
    final Set<Relationship> undefined = new HashSet<>();
    final Set<Relationship> relationships;
    final Processor processor = processorRef.get().getProcessor();
    try (final NarCloseable narCloseable = NarCloseable.withComponentNarLoader(processor.getClass(), processor.getIdentifier())) {
        relationships = processor.getRelationships();
    }
    if (relationships == null) {
        return undefined;
    }
    for (final Relationship relation : relationships) {
        final Set<Connection> connectionSet = this.connections.get(relation);
        if (connectionSet == null || connectionSet.isEmpty()) {
            undefined.add(relation);
        }
    }
    return undefined;
}

Example 79 with Connection

use of org.apache.nifi.connectable.Connection in project nifi by apache.

the class StandardProcessorNode method getIncomingNonLoopConnections.

List<Connection> getIncomingNonLoopConnections() {
    final List<Connection> connections = getIncomingConnections();
    final List<Connection> nonLoopConnections = new ArrayList<>(connections.size());
    for (final Connection connection : connections) {
        if (!connection.getSource().equals(this)) {
            nonLoopConnections.add(connection);
        }
    }
    return nonLoopConnections;
}

Example 80 with Connection

use of org.apache.nifi.connectable.Connection in project nifi by apache.

the class RepositoryContext method isAnyRelationshipAvailable.

public boolean isAnyRelationshipAvailable() {
    for (final Relationship relationship : getConnectable().getRelationships()) {
        final Collection<Connection> connections = getConnections(relationship);
        boolean available = true;
        for (final Connection connection : connections) {
            if (connection.getFlowFileQueue().isFull()) {
                available = false;
                break;
            }
        }
        if (available) {
            return true;
        }
    }
    return false;
}