NIFI Custom Processor Series Part 3: Junit Test with Nifi Mock

This discussion is a continuation of a NIFI Custom Processor Series Part 2. In Part 1 , we learned how to develop a custom processor from scratch and in Part 2, we improved that processor to support custom validators.

NOTE: If you are new to Apache NIFI, I would suggest you to check the official documentation of Apache NIFI. If you want a concise discussion covering the basics concepts of Apache NIFI, you can refer Chapter 8 of my book — Scalable Data Architecture with Java.

In the last part of our series, we developed a processor called SalutationProcessorv2 with a custom validator for property named — BeforeOrAfter . This processor produces a dynamic salutation message based on the incoming flowfile content. For example, Suppose the content of incoming flowfile is Sinchan, property value of Salutation property is set to Good Morning and BeforeOrAfter value is B, then output flowfile should have a content like Good Morning Sinchan.

In this part, we will learn how to write a junit test case to do unit testing of this processor called SalutationProcessorv2. We will use nifi-mock dependencies to create the junit test case. Now, let us see how we can add a junit test case for this processor

Step 1: Add the NIFI mock maven dependencies

In order to mock processing of a flowfile using NIFI processor, we need to add nifi-mock maven dpendencies in the pom.xml. However, since, it will be used only during test, so we will make the scope of the dependency as test. We will add this required dependency in pom.xml as follows

 <!-- https://mvnrepository.com/artifact/org.apache.nifi/nifi-mock -->
        <dependency>
            <groupId>org.apache.nifi</groupId>
            <artifactId>nifi-mock</artifactId>
            <version>${nifi.version}</version>
            <scope>test</scope>
        </dependency>

Next, Let’s create the test suite.

Step 2: Create a Test class and add test case

First, we should create a package with same name in test/java folder as that of the package of the processor we want to test.Under that package, we can create a Test class. Here, in our example, we have named the Test class as SalutationProcessorV2Test. The following screenshot shows how the project structure will look like

Figure — Screenshot of project structure after adding test class

Since we have created the test class, let us add the test case in the code. But before that, let us understand what use case we want to test. We want to test, how a salutation will look like if it is to be added before the name, coming from the flowfile. In order to do so, first we will add a method called testProcessorForSalutationBefore and annotate it with ‘@Test’ to make this a test case. The code block for the same is as shown below

public class SalutationProcessorV2Test {

    @Test
    public void testProcessorForSalutationBefore(){
       ...
    }
}

Next, we will be adding body of the method. First, we will create a nifi-mock’s TestRunner instance to mock a test run of SalutationProcessorV2 processor. This can be done as shown in the following code block

@Test
    public void testProcessorForSalutationBefore(){
        TestRunner testRunner = TestRunners.newTestRunner(new SalutationProcessorV2());
        ...
    }

Now, we will set the processor properties using the TestRunner object as follows

 @Test
    public void testProcessorForSalutationBefore(){
        ...
        //set the properties of the processor
        testRunner.setProperty(SalutationProcessorV2.SALUTATION,"Good Morning!");
        testRunner.setProperty(SalutationProcessorV2.BEFORE_OR_AFTER,"B");

        ...
    }

Then we will enqueue a flowfile using TestRunner.enqueue property. We can enqueue any number of messages using this API. We can enqueue using FlowFile object or by string content of the Flowfile. Here we will enqueue using flowfile content as follows

 @Test
    public void testProcessorForSalutationBefore(){
        ...
        //enqueue a flowfile
        testRunner.enqueue("Sinchan");

        ...
    }

Next, we will run the processor using TestRunner.run() method. It optionally takes an integer argument to indicate how many times we want the processor to run. In our case, we use run(1), which is similar to NIFI UI’s ‘Run Once’ functionality. Once it is run successfully, we check whether the flowfile in the input queue has been processed or not using TestRunner.assertQueueEmpty() method. The following code depicts the same

@Test
    public void testProcessorForSalutationBefore(){
        ...
       
        //execute 1 run - similar to NIFI UI's run once
        testRunner.run(1);

        //assert the input Q is empty and the flowfile is processed
        testRunner.assertQueueEmpty();

        ...
    }

Then, we list all the flowfiles from the Success queue of the SalutationProcessorV2 processor using the Testrunner method getFlowFilesForRelationship() as follows

@Test
    public void testProcessorForSalutationBefore(){
        ...

        //get the flowfiles from Success Q
        List<MockFlowFile> results = testRunner.getFlowFilesForRelationship(SalutationProcessorV2.SUCCESS);
        assertTrue("Number of flowfiles in Success Queue is as expected (No of Flowfile = 1) ",results.size()==1);

        ...
    }

As you can see in the above code block, we pass the relationship reference object (SalutationProcessorV2.SUCCESS) to getFlowFilesForRelationship() method, to obtain a List of output flowfiles ( which are of Type MockFlowFile as it is generated using nifi-mock APIs).Note that after the flowfiles are listed, we assert that only one flowfile is present in the success queue by using static method org.junit.Assert.assertTrue. Finally, we get the flowfile and extract the content of the flowfile and compare it with expected result as follows:

@Test
    public void testProcessorForSalutationBefore(){
        ...

        MockFlowFile outputFlowfile = results.get(0);
        String outputFlowfileContent = new String(testRunner.getContentAsByteArray(outputFlowfile));
        assertEquals("Good Morning!  Sinchan",outputFlowfileContent);
    }

As shown in the code block above, we store the content of the flowfile in outputFlowfileContent variable by extracting the content using TestRunner’s utility method getContentAsByteArray and converting the byte array to a String object. Finally, we compare the expected content with the actual output flowfile content using the static method — org.junit.Assert.assertEquals.

Now, that we have developed our test suite to test NIFI processor, let us test it.

Step 3: Run and verify Junit test

You can run the Junit test by either running the following maven command

mvn test

Or you can use the built-in feature of an IDE like Jetbrain’s IntelliJ to execute the junit test case. In intelliJ, you can right click on a test case and click ‘Run <name of the test case method>’ as shown in the screen shot below

Figure — Running Junit test from IntelliJ

Once the junit test are run, we get the test case output (passed or failed) along with console logs or errors as shown in the screenshot below

Figure — Ouput of Junit test case of NIFI processor

As seen in this example, we were able to successfully create and test a junit test case for the custom NIFI processor — SalutationProcessorV2 . You can find the full code base for this blog at: https://github.com/perfecxus/SalutationProcessor/blob/master/src/test/java/org/example/SalutationProcessorV2Test.java

Hope you found this read useful. Please feel free to discuss or post question in the comment section below.

I will see you again in the Part 4 of the blog series. Till then Happy Learning and Stay Safe !!