nifi-nar-bundles/nifi-splunk-bundle/nifi-splunk-processors/src/test/java/org/apache/nifi/processors/splunk/TestPutSplunk.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.nifi.processors.splunk;

import org.apache.nifi.logging.ComponentLog;
import org.apache.nifi.processor.util.put.sender.ChannelSender;
import org.apache.nifi.util.MockFlowFile;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import javax.net.ssl.SSLContext;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


public class TestPutSplunk {

    private TestRunner runner;
    private TestablePutSplunk proc;
    private CapturingChannelSender sender;

    @Before
    public void init() {
        ComponentLog logger = Mockito.mock(ComponentLog.class);
        sender = new CapturingChannelSender("localhost", 12345, 0, logger);
        proc = new TestablePutSplunk(sender);

        runner = TestRunners.newTestRunner(proc);
        runner.setProperty(PutSplunk.PORT, "12345");
    }

    @Test
    public void testUDPSendWholeFlowFile() {
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.UDP_VALUE.getValue());
        final String message = "This is one message, should send the whole FlowFile";

        runner.enqueue(message);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(1, sender.getMessages().size());
        Assert.assertEquals(message, sender.getMessages().get(0));
    }

    @Test
    public void testTCPSendWholeFlowFile() {
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        final String message = "This is one message, should send the whole FlowFile";

        runner.enqueue(message);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(1, sender.getMessages().size());
        Assert.assertEquals(message + "\n", sender.getMessages().get(0));
    }

    @Test
    public void testTCPSendWholeFlowFileAlreadyHasNewLine() {
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        final String message = "This is one message, should send the whole FlowFile\n";

        runner.enqueue(message);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(1, sender.getMessages().size());
        Assert.assertEquals(message, sender.getMessages().get(0));
    }

    @Test
    public void testUDPSendDelimitedMessages() {
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.UDP_VALUE.getValue());

        final String delimiter = "DD";
        runner.setProperty(PutSplunk.MESSAGE_DELIMITER, delimiter);

        final String message = "This is message 1DDThis is message 2DDThis is message 3";

        runner.enqueue(message);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(3, sender.getMessages().size());
        Assert.assertEquals("This is message 1", sender.getMessages().get(0));
        Assert.assertEquals("This is message 2", sender.getMessages().get(1));
        Assert.assertEquals("This is message 3", sender.getMessages().get(2));
    }

    @Test
    public void testTCPSendDelimitedMessages() {
        final String delimiter = "DD";
        runner.setProperty(PutSplunk.MESSAGE_DELIMITER, delimiter);
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        // no delimiter at end
        final String message = "This is message 1DDThis is message 2DDThis is message 3";

        runner.enqueue(message);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(3, sender.getMessages().size());
        Assert.assertEquals("This is message 1\n", sender.getMessages().get(0));
        Assert.assertEquals("This is message 2\n", sender.getMessages().get(1));
        Assert.assertEquals("This is message 3\n", sender.getMessages().get(2));
    }

    @Test
    public void testTCPSendDelimitedMessagesWithEL() {
        final String delimiter = "DD";
        runner.setProperty(PutSplunk.MESSAGE_DELIMITER, "${flow.file.delim}");
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        // no delimiter at end
        final String message = "This is message 1DDThis is message 2DDThis is message 3";

        final Map<String,String> attrs = new HashMap<>();
        attrs.put("flow.file.delim", delimiter);

        runner.enqueue(message, attrs);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(3, sender.getMessages().size());
        Assert.assertEquals("This is message 1\n", sender.getMessages().get(0));
        Assert.assertEquals("This is message 2\n", sender.getMessages().get(1));
        Assert.assertEquals("This is message 3\n", sender.getMessages().get(2));
    }

    @Test
    public void testTCPSendDelimitedMessagesEndsWithDelimiter() {
        final String delimiter = "DD";
        runner.setProperty(PutSplunk.MESSAGE_DELIMITER, delimiter);
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        // delimiter at end
        final String message = "This is message 1DDThis is message 2DDThis is message 3DD";

        runner.enqueue(message);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(3, sender.getMessages().size());
        Assert.assertEquals("This is message 1\n", sender.getMessages().get(0));
        Assert.assertEquals("This is message 2\n", sender.getMessages().get(1));
        Assert.assertEquals("This is message 3\n", sender.getMessages().get(2));
    }

    @Test
    public void testTCPSendDelimitedMessagesWithNewLineDelimiter() {
        final String delimiter = "\\n";
        runner.setProperty(PutSplunk.MESSAGE_DELIMITER, delimiter);
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        final String message = "This is message 1\nThis is message 2\nThis is message 3";

        runner.enqueue(message);
        runner.run(1);
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(3, sender.getMessages().size());
        Assert.assertEquals("This is message 1\n", sender.getMessages().get(0));
        Assert.assertEquals("This is message 2\n", sender.getMessages().get(1));
        Assert.assertEquals("This is message 3\n", sender.getMessages().get(2));
    }

    @Test
    public void testTCPSendDelimitedMessagesWithErrors() {
        sender.setErrorStart(3);
        sender.setErrorEnd(4);

        final String delimiter = "DD";
        runner.setProperty(PutSplunk.MESSAGE_DELIMITER, delimiter);
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        // no delimiter at end
        final String success = "This is message 1DDThis is message 2DD";
        final String failure = "This is message 3DDThis is message 4";
        final String message = success + failure;

        runner.enqueue(message);
        runner.run(1);
        runner.assertTransferCount(PutSplunk.REL_SUCCESS, 1);
        runner.assertTransferCount(PutSplunk.REL_FAILURE, 1);

        // first two messages should went out success
        final MockFlowFile successFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        successFlowFile.assertContentEquals(success);

        // second two messages should went to failure
        final MockFlowFile failureFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_FAILURE).get(0);
        failureFlowFile.assertContentEquals(failure);

        // should only have the first two messages
        Assert.assertEquals(2, sender.getMessages().size());
        Assert.assertEquals("This is message 1\n", sender.getMessages().get(0));
        Assert.assertEquals("This is message 2\n", sender.getMessages().get(1));
    }

    @Test
    public void testTCPSendDelimitedMessagesWithErrorsInMiddle() {
        sender.setErrorStart(3);
        sender.setErrorEnd(4);

        final String delimiter = "DD";
        runner.setProperty(PutSplunk.MESSAGE_DELIMITER, delimiter);
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.TCP_VALUE.getValue());

        // no delimiter at end
        final String success = "This is message 1DDThis is message 2DD";
        final String failure = "This is message 3DDThis is message 4DD";
        final String success2 = "This is message 5DDThis is message 6DDThis is message 7DD";
        final String message = success + failure + success2;

        runner.enqueue(message);
        runner.run(1);
        runner.assertTransferCount(PutSplunk.REL_SUCCESS, 2);
        runner.assertTransferCount(PutSplunk.REL_FAILURE, 1);

        // first two messages should have went out success
        final MockFlowFile successFlowFile1 = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        successFlowFile1.assertContentEquals(success);

        // last three messages should have went out success
        final MockFlowFile successFlowFile2 = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(1);
        successFlowFile2.assertContentEquals(success2);

        // second two messages should have went to failure
        final MockFlowFile failureFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_FAILURE).get(0);
        failureFlowFile.assertContentEquals(failure);

        // should only have the first two messages
        Assert.assertEquals(5, sender.getMessages().size());
        Assert.assertEquals("This is message 1\n", sender.getMessages().get(0));
        Assert.assertEquals("This is message 2\n", sender.getMessages().get(1));
        Assert.assertEquals("This is message 5\n", sender.getMessages().get(2));
        Assert.assertEquals("This is message 6\n", sender.getMessages().get(3));
        Assert.assertEquals("This is message 7\n", sender.getMessages().get(4));
    }

    @Test
    public void testCompletingPreviousBatchOnNextExecution() {
        runner.setProperty(PutSplunk.PROTOCOL, PutSplunk.UDP_VALUE.getValue());

        final String message = "This is one message, should send the whole FlowFile";

        runner.enqueue(message);
        runner.run(2, false); // don't shutdown to prove that next onTrigger complete previous batch
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_SUCCESS, 1);

        final MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(PutSplunk.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(message);

        Assert.assertEquals(1, sender.getMessages().size());
        Assert.assertEquals(message, sender.getMessages().get(0));
    }

    @Test
    public void testUnableToCreateConnectionShouldRouteToFailure() {
        PutSplunk proc = new UnableToConnectPutSplunk();
        runner = TestRunners.newTestRunner(proc);
        runner.setProperty(PutSplunk.PORT, "12345");

        final String message = "This is one message, should send the whole FlowFile";

        runner.enqueue(message);
        runner.run();
        runner.assertAllFlowFilesTransferred(PutSplunk.REL_FAILURE, 1);
    }

    /**
     * Extend PutSplunk to use a CapturingChannelSender.
     */
    private static class UnableToConnectPutSplunk extends PutSplunk {

        @Override
        protected ChannelSender createSender(String protocol, String host, int port, int timeout, int maxSendBufferSize, SSLContext sslContext) throws IOException {
            throw new IOException("Unable to create connection");
        }
    }

    /**
     * Extend PutSplunk to use a CapturingChannelSender.
     */
    private static class TestablePutSplunk extends PutSplunk {

        private ChannelSender sender;

        public TestablePutSplunk(ChannelSender channelSender) {
            this.sender = channelSender;
        }

        @Override
        protected ChannelSender createSender(String protocol, String host, int port, int timeout, int maxSendBufferSize, SSLContext sslContext) throws IOException {
            return sender;
        }
    }


    /**
     * A ChannelSender that captures each message that was sent.
     */
    private static class CapturingChannelSender extends ChannelSender {

        private List<String> messages = new ArrayList<>();
        private int count = 0;
        private int errorStart = -1;
        private int errorEnd = -1;

        public CapturingChannelSender(String host, int port, int maxSendBufferSize, ComponentLog logger) {
            super(host, port, maxSendBufferSize, logger);
        }

        @Override
        public void open() throws IOException {

        }

        @Override
        protected void write(byte[] data) throws IOException {
            count++;
            if (errorStart > 0 && count >= errorStart && errorEnd > 0 && count <= errorEnd) {
                throw new IOException("this is an error");
            }
            messages.add(new String(data, StandardCharsets.UTF_8));
        }

        @Override
        public boolean isConnected() {
            return false;
        }

        @Override
        public void close() {

        }

        public List<String> getMessages() {
            return messages;
        }

        public void setErrorStart(int errorStart) {
            this.errorStart = errorStart;
        }

        public void setErrorEnd(int errorEnd) {
            this.errorEnd = errorEnd;
        }
    }

}