org.apache.nifi.remote.TransferDirection Java Examples

@Test
public void testSendOneFlowFile() throws IOException {

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    DataOutputStream serverResponse = new DataOutputStream(serverResponseBos);
    ResponseCode.CONFIRM_TRANSACTION.writeResponse(serverResponse, "2946083981");
    ResponseCode.TRANSACTION_FINISHED.writeResponse(serverResponse);

    ByteArrayInputStream bis = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream bos = new ByteArrayOutputStream();

    SocketClientTransaction transaction = getClientTransaction(bis, bos, TransferDirection.SEND);

    execSendOneFlowFile(transaction);

    // Verify what client has sent.
    DataInputStream sentByClient = new DataInputStream(new ByteArrayInputStream(bos.toByteArray()));
    assertEquals(RequestType.SEND_FLOWFILES, RequestType.readRequestType(sentByClient));
    DataPacket packetByClient = codec.decode(sentByClient);
    assertEquals("contents on client 1", readContents(packetByClient));
    Response endOfDataResponse = Response.read(sentByClient);
    assertEquals(ResponseCode.FINISH_TRANSACTION, endOfDataResponse.getCode());
    Response confirmResponse = Response.read(sentByClient);
    assertEquals(ResponseCode.CONFIRM_TRANSACTION, confirmResponse.getCode());
    assertEquals(-1, sentByClient.read());
}
 

@Override
protected SiteToSiteClient getClient() {
    final SiteToSiteClient client = Mockito.mock(SiteToSiteClient.class);
    final Transaction transaction = Mockito.mock(Transaction.class);

    try {
        Mockito.doAnswer(new Answer<Object>() {
            @Override
            public Object answer(final InvocationOnMock invocation) throws Throwable {
                final byte[] data = invocation.getArgumentAt(0, byte[].class);
                dataSent.add(data);
                return null;
            }
        }).when(transaction).send(Mockito.any(byte[].class), Mockito.any(Map.class));

        Mockito.when(client.createTransaction(Mockito.any(TransferDirection.class))).thenReturn(transaction);
    } catch (final Exception e) {
        e.printStackTrace();
        Assert.fail(e.toString());
    }

    return client;
}
 

@Override
protected SiteToSiteClient getClient() {
    final SiteToSiteClient client = Mockito.mock(SiteToSiteClient.class);
    final Transaction transaction = Mockito.mock(Transaction.class);

    try {
        Mockito.doAnswer(new Answer<Object>() {
            @Override
            public Object answer(final InvocationOnMock invocation) throws Throwable {
                final byte[] data = invocation.getArgumentAt(0, byte[].class);
                dataSent.add(data);
                return null;
            }
        }).when(transaction).send(Mockito.any(byte[].class), Mockito.anyMapOf(String.class, String.class));

        Mockito.when(client.createTransaction(Mockito.any(TransferDirection.class))).thenReturn(transaction);
    } catch (final Exception e) {
        e.printStackTrace();
        Assert.fail(e.toString());
    }

    return client;
}
 

@Test
public void testReceiveTwoFlowFiles() throws IOException {

    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doReturn(true).when(apiClient).openConnectionForReceive(eq(transactionUrl), any(Peer.class));
    TransactionResultEntity resultEntity = new TransactionResultEntity();
    resultEntity.setResponseCode(CONFIRM_TRANSACTION.getCode());
    doReturn(resultEntity).when(apiClient).commitReceivingFlowFiles(eq(transactionUrl), eq(CONFIRM_TRANSACTION), eq("2969091230"));

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    codec.encode(createDataPacket("contents on server 1"), serverResponseBos);
    codec.encode(createDataPacket("contents on server 2"), serverResponseBos);
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.RECEIVE, transactionUrl);

    execReceiveTwoFlowFiles(transaction);

    assertEquals("Client sends nothing as payload to receive flow files.", 0, clientRequest.toByteArray().length);
    verify(apiClient).commitReceivingFlowFiles(transactionUrl, CONFIRM_TRANSACTION, "2969091230");
}
 

@Test
public void testReceiveTimeout() throws Exception {

    try (
            SiteToSiteClient client = getDefaultBuilder()
                    .timeout(1, TimeUnit.SECONDS)
                    .portName("output-timeout")
                    .build()
    ) {
        try {
            client.createTransaction(TransferDirection.RECEIVE);
            fail();
        } catch (IOException e) {
            logger.info("An exception was thrown as expected.", e);
            assertTrue(e instanceof SocketTimeoutException);
        }
    }
}
 

@Test
public void testSendZeroFlowFile() throws IOException {

    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doNothing().when(apiClient).openConnectionForSend(eq(transactionUrl), any(Peer.class));

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.SEND, transactionUrl);

    execSendZeroFlowFile(transaction);

    assertEquals("Client didn't send anything", 0, clientRequest.toByteArray().length);
}
 

@Test
public void testUnkownClusterUrl() throws Exception {

    final URI uri = server.getURI();

    try (
        SiteToSiteClient client = getDefaultBuilder()
            .url("http://" + uri.getHost() + ":" + uri.getPort() + "/unkown")
            .portName("input-running")
            .build()
    ) {
        final Transaction transaction = client.createTransaction(TransferDirection.SEND);

        assertNull(transaction);

    }

}
 

@Test
public void testSendZeroFlowFile() throws IOException {

    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doNothing().when(apiClient).openConnectionForSend(eq(transactionUrl), any(Peer.class));

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.SEND, transactionUrl);

    execSendZeroFlowFile(transaction);

    assertEquals("Client didn't send anything", 0, clientRequest.toByteArray().length);
}
 

@Test
public void testReceiveWithInvalidChecksum() throws IOException {

    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doReturn(true).when(apiClient).openConnectionForReceive(eq(transactionUrl), any(Peer.class));
    // The checksum is correct, but here we simulate as if it's wrong, BAD_CHECKSUM.
    TransactionResultEntity resultEntity = new TransactionResultEntity();
    resultEntity.setResponseCode(ResponseCode.BAD_CHECKSUM.getCode());
    doReturn(resultEntity).when(apiClient).commitReceivingFlowFiles(eq(transactionUrl), eq(CONFIRM_TRANSACTION), eq("2969091230"));

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    codec.encode(createDataPacket("contents on server 1"), serverResponseBos);
    codec.encode(createDataPacket("contents on server 2"), serverResponseBos);
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.RECEIVE, transactionUrl);

    execReceiveWithInvalidChecksum(transaction);

    assertEquals("Client sends nothing as payload to receive flow files.", 0, clientRequest.toByteArray().length);
    verify(apiClient).commitReceivingFlowFiles(transactionUrl, CONFIRM_TRANSACTION, "2969091230");
}
 

@Test
public void testFormulateDestinationListForOutput() throws IOException {
    final Set<PeerStatus> collection = new HashSet<>();
    collection.add(new PeerStatus(new PeerDescription("HasMedium", 1111, true), 4096, true));
    collection.add(new PeerStatus(new PeerDescription("HasLots", 2222, true), 10240, true));
    collection.add(new PeerStatus(new PeerDescription("HasLittle", 3333, true), 1024, true));
    collection.add(new PeerStatus(new PeerDescription("HasMedium", 4444, true), 4096, true));
    collection.add(new PeerStatus(new PeerDescription("HasMedium", 5555, true), 4096, true));

    PeerStatusProvider peerStatusProvider = Mockito.mock(PeerStatusProvider.class);
    PeerSelector peerSelector = new PeerSelector(peerStatusProvider, null);

    final List<PeerStatus> destinations = peerSelector.formulateDestinationList(collection, TransferDirection.RECEIVE);
    final Map<String, Integer> selectedCounts = calculateAverageSelectedCount(collection, destinations);

    logger.info("selectedCounts={}", selectedCounts);
    assertTrue("HasLots should send lots", selectedCounts.get("HasLots") > selectedCounts.get("HasMedium"));
    assertTrue("HasMedium should send medium", selectedCounts.get("HasMedium") > selectedCounts.get("HasLittle"));
}
 

@Test
public void testReceiveTwoFlowFiles() throws IOException {

    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doReturn(true).when(apiClient).openConnectionForReceive(eq(transactionUrl), any(Peer.class));
    TransactionResultEntity resultEntity = new TransactionResultEntity();
    resultEntity.setResponseCode(CONFIRM_TRANSACTION.getCode());
    doReturn(resultEntity).when(apiClient).commitReceivingFlowFiles(eq(transactionUrl), eq(CONFIRM_TRANSACTION), eq("2969091230"));

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    codec.encode(createDataPacket("contents on server 1"), serverResponseBos);
    codec.encode(createDataPacket("contents on server 2"), serverResponseBos);
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.RECEIVE, transactionUrl);

    execReceiveTwoFlowFiles(transaction);

    assertEquals("Client sends nothing as payload to receive flow files.", 0, clientRequest.toByteArray().length);
    verify(apiClient).commitReceivingFlowFiles(transactionUrl, CONFIRM_TRANSACTION, "2969091230");
}
 

@Before
public void setup() {
    SiteToSiteClient.Builder builder = new SiteToSiteClient.Builder();
    expectedUrl = SiteToSiteCliMain.URL_OPTION_DEFAULT;
    expectedTransferDirection = TransferDirection.valueOf(SiteToSiteCliMain.DIRECTION_OPTION_DEFAULT);
    expectedSiteToSiteTransportProtocol = SiteToSiteTransportProtocol.valueOf(SiteToSiteCliMain.TRANSPORT_PROTOCOL_OPTION_DEFAULT);
    expectedPortName = builder.getPortName();
    expectedPortIdentifier = builder.getPortIdentifier();
    expectedTimeoutNs = builder.getTimeout(TimeUnit.NANOSECONDS);
    expectedPenalizationNs = builder.getPenalizationPeriod(TimeUnit.NANOSECONDS);
    expectedKeystoreFilename = builder.getKeystoreFilename();
    expectedKeystorePass = builder.getKeystorePass();
    expectedKeystoreType = builder.getKeystoreType();
    expectedTruststoreFilename = builder.getTruststoreFilename();
    expectedTruststorePass = builder.getTruststorePass();
    expectedTruststoreType = builder.getTruststoreType();
    expectedCompression = false;
    expectedPeerPersistenceFile = builder.getPeerPersistenceFile();
    SiteToSiteClientConfig siteToSiteClientConfig = builder.buildConfig();
    expectedBatchCount = siteToSiteClientConfig.getPreferredBatchCount();
    expectedBatchDuration = siteToSiteClientConfig.getPreferredBatchDuration(TimeUnit.NANOSECONDS);
    expectedBatchSize = siteToSiteClientConfig.getPreferredBatchSize();
    expectedHttpProxy = siteToSiteClientConfig.getHttpProxy();
}
 

@Test
public void testReceiveTimeout() throws Exception {

    try (
            SiteToSiteClient client = getDefaultBuilder()
                    .timeout(1, TimeUnit.SECONDS)
                    .portName("output-timeout")
                    .build()
    ) {
        try {
            client.createTransaction(TransferDirection.RECEIVE);
            fail();
        } catch (IOException e) {
            logger.info("An exception was thrown as expected.", e);
            assertTrue(e instanceof SocketTimeoutException);
        }
    }
}
 

/**
 * Returns a map of peers prepared for flowfile transfer in the specified direction. Each peer is a key and the value is a
 * weighted percentage of the total flowfiles in the remote instance. For example, in a cluster where the total number of flowfiles
 * is 100, distributed across three nodes 20 in A, 30 in B, and 50 in C, the resulting map for
 * {@code SEND} will be {@code [A:40.0, B:35.0, C:25.0]} (1 - .2 => .8 * 100 / (3-1)) => 40.0).
 *
 * @param statuses  the set of all peers
 * @param direction the direction of transfer ({@code SEND} weights the destinations higher if they have more flowfiles, {@code RECEIVE} weights them higher if they have fewer)
 * @return the ordered map of each peer to its relative weight
 */
LinkedHashMap<PeerStatus, Double> buildWeightedPeerMap(final Set<PeerStatus> statuses, final TransferDirection direction) {
    // Get all the destinations with their relative weights
    final Map<PeerStatus, Double> peerWorkloads = createDestinationMap(statuses, direction);

    if (!peerWorkloads.isEmpty()) {
        // This map is sorted, but not by key, so it cannot use SortedMap
        LinkedHashMap<PeerStatus, Double> sortedPeerWorkloads = sortMapByWeight(peerWorkloads);

        // Print the expected distribution of the peers
        printDistributionStatistics(sortedPeerWorkloads, direction);

        return sortedPeerWorkloads;
    } else {
        logger.debug("No peers available");
        return new LinkedHashMap<>();
    }
}
 

private String getPortIdentifier(final TransferDirection direction) throws IOException {
    final String id = this.portIdentifier;
    if (id != null) {
        return id;
    }

    final String portId;
    if (direction == TransferDirection.SEND) {
        portId = siteInfoProvider.getInputPortIdentifier(this.portName);
    } else {
        portId = siteInfoProvider.getOutputPortIdentifier(this.portName);
    }

    if (portId == null) {
        logger.debug("Unable to resolve port [{}] to an identifier", portName);
    } else {
        logger.debug("Resolved port [{}] to identifier [{}]", portName, portId);
        this.portIdentifier = portId;
    }

    return portId;
}
 

@Test
public void testSendZeroFlowFile() throws IOException {

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();

    ByteArrayInputStream bis = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream bos = new ByteArrayOutputStream();

    SocketClientTransaction transaction = getClientTransaction(bis, bos, TransferDirection.SEND);

    execSendZeroFlowFile(transaction);

    // Verify what client has sent.
    DataInputStream sentByClient = new DataInputStream(new ByteArrayInputStream(bos.toByteArray()));
    assertEquals(RequestType.SEND_FLOWFILES, RequestType.readRequestType(sentByClient));
    assertEquals(-1, sentByClient.read());
}
 

@Override
public boolean runRecursive(final Queue<InMemoryFlowFile> queue) {
    try {
        final Transaction transaction = client.createTransaction(TransferDirection.SEND);
        if (transaction == null) {
            getLogger().error("Unable to create a transaction for Remote Process Group {} to send to port {}", new Object[] {url, name});
            return false;
        }

        StatelessFlowFile flowFile;
        while ((flowFile = inputQueue.poll()) != null) {
            final DataPacket dataPacket = new StandardDataPacket(flowFile.getAttributes(), flowFile.getDataStream(), flowFile.getSize());
            transaction.send(dataPacket);
        }

        transaction.confirm();
        transaction.complete();
    } catch (final Exception e) {
        getLogger().error("Failed to send FlowFile via site-to-site", e);
        return false;
    }

    return true;
}
 

@Test
public void testReceiveSlowClientSuccess() throws Exception {

    try (
            SiteToSiteClient client = getDefaultBuilder()
                    .portName("output-running")
                    .build()
    ) {
        final Transaction transaction = client.createTransaction(TransferDirection.RECEIVE);

        assertNotNull(transaction);

        DataPacket packet;
        while ((packet = transaction.receive()) != null) {
            consumeDataPacket(packet);
            Thread.sleep(500);
        }
        transaction.confirm();
        transaction.complete();
    }
}
 

@Test
public void testReceiveUnknownPort() throws Exception {

    try (
        SiteToSiteClient client = getDefaultBuilder()
            .portName("output-unknown")
            .build()
    ) {
        try {
            client.createTransaction(TransferDirection.RECEIVE);
            fail();
        } catch (IOException e) {
            logger.info("Exception message: {}", e.getMessage());
            assertTrue(e.getMessage().contains("Failed to determine the identifier of port"));
        }
    }
}
 

@Test
public void testReceiveZeroFlowFile() throws IOException {

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    DataOutputStream serverResponse = new DataOutputStream(serverResponseBos);
    ResponseCode.NO_MORE_DATA.writeResponse(serverResponse);

    ByteArrayInputStream bis = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream bos = new ByteArrayOutputStream();

    SocketClientTransaction transaction = getClientTransaction(bis, bos, TransferDirection.RECEIVE);

    execReceiveZeroFlowFile(transaction);

    // Verify what client has sent.
    DataInputStream sentByClient = new DataInputStream(new ByteArrayInputStream(bos.toByteArray()));
    assertEquals(RequestType.RECEIVE_FLOWFILES, RequestType.readRequestType(sentByClient));
    assertEquals(-1, sentByClient.read());
}
 

@Override
public void setup(ReportingContext reportContext) throws IOException {
    if(siteToSiteClient == null) {
        final SiteToSiteClient client = Mockito.mock(SiteToSiteClient.class);
        final Transaction transaction = Mockito.mock(Transaction.class);

        try {
            Mockito.doAnswer((Answer<Object>) invocation -> {
                final byte[] data = invocation.getArgument(0, byte[].class);
                dataSent.add(data);
                return null;
            }).when(transaction).send(Mockito.any(byte[].class), Mockito.any(Map.class));

            when(client.createTransaction(Mockito.any(TransferDirection.class))).thenReturn(transaction);
        } catch (final Exception e) {
            e.printStackTrace();
            Assert.fail(e.toString());
        }
        siteToSiteClient = client;
    }
}
 

private SocketClientTransaction getClientTransaction(ByteArrayInputStream bis, ByteArrayOutputStream bos, TransferDirection direction) throws IOException {
    PeerDescription description = null;
    String peerUrl = "";
    SocketCommunicationsSession commsSession = mock(SocketCommunicationsSession.class);
    SocketInput socketIn = mock(SocketInput.class);
    SocketOutput socketOut = mock(SocketOutput.class);
    when(commsSession.getInput()).thenReturn(socketIn);
    when(commsSession.getOutput()).thenReturn(socketOut);

    when(socketIn.getInputStream()).thenReturn(bis);
    when(socketOut.getOutputStream()).thenReturn(bos);

    String clusterUrl = "";
    Peer peer = new Peer(description, commsSession, peerUrl, clusterUrl);
    boolean useCompression = false;
    int penaltyMillis = 1000;
    EventReporter eventReporter = null;
    int protocolVersion = 5;
    String destinationId = "destinationId";
    return new SocketClientTransaction(protocolVersion, destinationId, peer, codec, direction, useCompression, penaltyMillis, eventReporter);
}
 

@Test
public void testNoAvailablePeer() throws Exception {

    peers = new HashSet<>();

    try (
        SiteToSiteClient client = getDefaultBuilder()
            .portName("input-running")
            .build()
    ) {
        final Transaction transaction = client.createTransaction(TransferDirection.SEND);

        assertNull(transaction);

    }

}
 

/**
 * Returns the normalized weight for this ratio of peer flowfiles to total flowfiles and the given direction. The number will be
 * a Double between 0 and 100 indicating the percent of all flowfiles the peer
 * should send/receive. The transfer direction is <em>from the perspective of this node to the peer</em>
 * (i.e. how many flowfiles should <em>this node send</em> to the peer, or how many flowfiles
 * should <em>this node receive</em> from the peer).
 *
 * @param direction          the transfer direction ({@code SEND} weights the destinations higher if they have fewer flowfiles, {@code RECEIVE} weights them higher if they have more)
 * @param totalFlowFileCount the total flowfile count in the remote instance (standalone or cluster)
 * @param flowFileCount      the flowfile count for the given peer
 * @param peerCount          the number of peers in the remote instance
 * @return the normalized weight of this peer
 */
private static double calculateNormalizedWeight(TransferDirection direction, long totalFlowFileCount, int flowFileCount, int peerCount) {
    // If there is only a single remote, send/receive all data to/from it
    if (peerCount == 1) {
        return 100;
    }

    double cappedPercent;
    // If no flowfiles exist in the remote instance, evenly weight each node with 1/N
    if (totalFlowFileCount == 0) {
        cappedPercent = 1.0 / peerCount;
    } else {
        final double percentageOfFlowFiles = ((double) flowFileCount / totalFlowFileCount);
        cappedPercent = percentageOfFlowFiles;

        // If sending to the remote, allocate more flowfiles to the less-stressed peers
        if (direction == TransferDirection.SEND) {
            cappedPercent = (1 - percentageOfFlowFiles) / (peerCount - 1);
        }
    }
    return new BigDecimal(cappedPercent * 100).setScale(2, RoundingMode.FLOOR).doubleValue();
}
 

@Test
public void testSendZeroFlowFile() throws IOException {

    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();

    ByteArrayInputStream bis = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream bos = new ByteArrayOutputStream();

    SocketClientTransaction transaction = getClientTransaction(bis, bos, TransferDirection.SEND);

    execSendZeroFlowFile(transaction);

    // Verify what client has sent.
    DataInputStream sentByClient = new DataInputStream(new ByteArrayInputStream(bos.toByteArray()));
    assertEquals(RequestType.SEND_FLOWFILES, RequestType.readRequestType(sentByClient));
    assertEquals(-1, sentByClient.read());
}
 

@Test
public void testSendTimeout() throws Exception {

    try (
        SiteToSiteClient client = getDefaultBuilder()
            .timeout(1, TimeUnit.SECONDS)
            .portName("input-timeout")
            .build()
    ) {
        final Transaction transaction = client.createTransaction(TransferDirection.SEND);

        assertNotNull(transaction);

        DataPacket packet = new DataPacketBuilder()
            .contents("Example contents from client.")
            .attr("Client attr 1", "Client attr 1 value")
            .attr("Client attr 2", "Client attr 2 value")
            .build();
        serverChecksum = "1345413116";

        transaction.send(packet);
        try {
            transaction.confirm();
            fail();
        } catch (IOException e) {
            logger.info("An exception was thrown as expected.", e);
            assertTrue(e.getMessage().contains("TimeoutException"));
        }

        completeShouldFail(transaction);
    }

}
 

private static void testSendIgnoreProxyError(final SiteToSiteClient client, final SendData function) throws IOException {
    final boolean isProxyEnabled = client.getConfig().getHttpProxy() != null;
    try {
        final Transaction transaction = client.createTransaction(TransferDirection.SEND);

        if (isProxyEnabled && transaction == null) {
            // Transaction is not created sometimes at AppVeyor.
            logger.warn("Transaction was not created. Most likely an environment dependent issue.");
            return;
        }

        assertNotNull(transaction);

        function.apply(transaction);

        transaction.confirm();

        transaction.complete();
    } catch (final IOException e) {
        if (isProxyEnabled && e.getMessage().contains("504")) {
            // Gateway Timeout happens sometimes at Travis CI.
            logger.warn("Request timeout. Most likely an environment dependent issue.", e);
        } else {
            throw e;
        }
    }
}
 

@Test
@Ignore("For local testing only; not really a unit test but a manual test")
public void testSend() throws IOException {
    System.setProperty("org.slf4j.simpleLogger.log.org.apache.nifi.remote", "DEBUG");

    final SiteToSiteClient client = new SiteToSiteClient.Builder()
            .url("http://localhost:8080/nifi")
            .portName("input")
            .build();

    try {
        final Transaction transaction = client.createTransaction(TransferDirection.SEND);
        Assert.assertNotNull(transaction);

        final Map<String, String> attrs = new HashMap<>();
        attrs.put("site-to-site", "yes, please!");
        final byte[] bytes = "Hello".getBytes();
        final ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
        final DataPacket packet = new StandardDataPacket(attrs, bais, bytes.length);
        transaction.send(packet);

        transaction.confirm();
        transaction.complete();
    } finally {
        client.close();
    }
}
 

@Test
@Ignore("For local testing only; not really a unit test but a manual test")
public void testReceive() throws IOException {
    System.setProperty("org.slf4j.simpleLogger.log.org.apache.nifi.remote", "DEBUG");

    final SiteToSiteClient client = new SiteToSiteClient.Builder()
            .url("http://localhost:8080/nifi")
            .portName("cba")
            .requestBatchCount(10)
            .build();

    try {
        for (int i = 0; i < 1000; i++) {
            final Transaction transaction = client.createTransaction(TransferDirection.RECEIVE);
            Assert.assertNotNull(transaction);

            DataPacket packet;
            while (true) {
                packet = transaction.receive();
                if (packet == null) {
                    break;
                }

                final InputStream in = packet.getData();
                final long size = packet.getSize();
                final byte[] buff = new byte[(int) size];

                StreamUtils.fillBuffer(in, buff);
            }

            transaction.confirm();
            transaction.complete();
        }
    } finally {
        client.close();
    }
}
 

@Test
public void testFormulateDestinationListForInputPortsHugeDifference() throws IOException {
    final Set<PeerStatus> collection = new HashSet<>();
    collection.add(new PeerStatus(new PeerDescription("HasLots", 1111, true), 500, true));
    collection.add(new PeerStatus(new PeerDescription("HasLittle", 2222, true), 50000, true));

    PeerStatusProvider peerStatusProvider = Mockito.mock(PeerStatusProvider.class);
    PeerSelector peerSelector = new PeerSelector(peerStatusProvider, null);

    final List<PeerStatus> destinations = peerSelector.formulateDestinationList(collection, TransferDirection.RECEIVE);
    final Map<String, Integer> selectedCounts = calculateAverageSelectedCount(collection, destinations);

    logger.info("selectedCounts={}", selectedCounts);
    assertTrue("HasLots should get little", selectedCounts.get("HasLots") < selectedCounts.get("HasLittle"));
}