Java tutorial
/* * 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.elasticsearch; import static org.junit.Assert.assertNotNull; import static org.mockito.Matchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import java.io.IOException; import java.util.Arrays; import java.util.HashMap; import java.util.List; import org.apache.commons.io.IOUtils; import org.apache.commons.lang3.StringUtils; import org.apache.nifi.processor.ProcessContext; import org.apache.nifi.processor.exception.ProcessException; import org.apache.nifi.ssl.SSLContextService; import org.apache.nifi.util.MockFlowFile; import org.apache.nifi.util.TestRunner; import org.apache.nifi.util.TestRunners; import org.junit.After; import org.junit.Assert; import org.junit.Test; import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; import org.mockito.stubbing.OngoingStubbing; import okhttp3.Call; import okhttp3.MediaType; import okhttp3.OkHttpClient; import okhttp3.Protocol; import okhttp3.Request; import okhttp3.Response; import okhttp3.ResponseBody; public class TestScrollElasticsearchHttp { private TestRunner runner; @After public void teardown() { runner = null; } @Test public void testScrollElasticsearchOnTrigger_withNoInput() throws IOException { runner = TestRunners.newTestRunner(new ScrollElasticsearchHttpTestProcessor()); runner.setValidateExpressionUsage(true); runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.assertNotValid(); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.assertNotValid(); runner.setProperty(ScrollElasticsearchHttp.QUERY, "source:WZ AND identifier:\"${identifier}\""); runner.assertValid(); runner.setProperty(ScrollElasticsearchHttp.PAGE_SIZE, "2"); runner.assertValid(); runner.setIncomingConnection(false); runAndVerifySuccess(); } @Test public void testScrollElasticsearchOnTrigger_withNoInput_EL() throws IOException { runner = TestRunners.newTestRunner(new ScrollElasticsearchHttpTestProcessor()); runner.setValidateExpressionUsage(true); runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "${es.url}"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.assertNotValid(); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.assertNotValid(); runner.setProperty(ScrollElasticsearchHttp.QUERY, "source:WZ AND identifier:\"${identifier}\""); runner.assertValid(); runner.setProperty(ScrollElasticsearchHttp.PAGE_SIZE, "2"); runner.assertValid(); runner.setProperty(AbstractElasticsearchHttpProcessor.CONNECT_TIMEOUT, "${connect.timeout}"); runner.assertValid(); runner.setVariable("es.url", "http://127.0.0.1:9200"); runner.setIncomingConnection(false); runAndVerifySuccess(); } private void runAndVerifySuccess() { runner.enqueue("".getBytes(), new HashMap<String, String>() { { put("identifier", "28039652140"); } }); // Must run once for each of the 3 pages runner.run(3, true, true); runner.assertAllFlowFilesTransferred(ScrollElasticsearchHttp.REL_SUCCESS, 2); final MockFlowFile out = runner.getFlowFilesForRelationship(ScrollElasticsearchHttp.REL_SUCCESS).get(0); assertNotNull(out); int numHits = runner.getFlowFilesForRelationship(ScrollElasticsearchHttp.REL_SUCCESS).stream().map(ff -> { String page = new String(ff.toByteArray()); return StringUtils.countMatches(page, "{\"timestamp\""); }).reduce((a, b) -> a + b).get(); Assert.assertEquals(3, numHits); } @Test public void testScrollElasticsearchOnTriggerWithFields() throws IOException { runner = TestRunners.newTestRunner(new ScrollElasticsearchHttpTestProcessor()); runner.setValidateExpressionUsage(true); runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.assertNotValid(); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.assertNotValid(); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.assertValid(); runner.setProperty(ScrollElasticsearchHttp.FIELDS, "id,, userinfo.location"); runner.assertValid(); runner.setProperty(ScrollElasticsearchHttp.SORT, "timestamp:asc,identifier:desc"); runner.assertValid(); runner.setIncomingConnection(false); runAndVerifySuccess(); } @Test public void testScrollElasticsearchOnTriggerWithServerFail() throws IOException { ScrollElasticsearchHttpTestProcessor processor = new ScrollElasticsearchHttpTestProcessor(); processor.setStatus(100, "Should fail"); runner = TestRunners.newTestRunner(processor); // simulate doc not found runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.setValidateExpressionUsage(true); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.setIncomingConnection(false); runner.enqueue("".getBytes(), new HashMap<String, String>() { { put("identifier", "28039652140"); } }); runner.run(1, true, true); // This test generates a HTTP 100 "Should fail" runner.assertTransferCount(ScrollElasticsearchHttp.REL_FAILURE, 0); runner.assertTransferCount(ScrollElasticsearchHttp.REL_SUCCESS, 0); } @Test public void testScrollElasticsearchOnTriggerWithServerRetry() throws IOException { ScrollElasticsearchHttpTestProcessor processor = new ScrollElasticsearchHttpTestProcessor(); processor.setStatus(500, "Internal error"); runner = TestRunners.newTestRunner(processor); // simulate doc not found runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.setValidateExpressionUsage(true); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.setIncomingConnection(false); runner.enqueue("".getBytes(), new HashMap<String, String>() { { put("identifier", "28039652140"); } }); runner.run(1, true, true); // This test generates a HTTP 500 "Internal error" runner.assertTransferCount(ScrollElasticsearchHttp.REL_FAILURE, 0); runner.assertTransferCount(ScrollElasticsearchHttp.REL_SUCCESS, 0); } @Test public void testScrollElasticsearchOnTriggerWithServerFailAfterSuccess() throws IOException { ScrollElasticsearchHttpTestProcessor processor = new ScrollElasticsearchHttpTestProcessor(); processor.setStatus(100, "Should fail", 2); runner = TestRunners.newTestRunner(processor); // simulate doc not found runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.setValidateExpressionUsage(true); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.enqueue("".getBytes(), new HashMap<String, String>() { { put("identifier", "28039652140"); } }); runner.setIncomingConnection(false); runner.run(1, true, true); // This test generates a HTTP 100 "Should fail" runner.assertTransferCount(ScrollElasticsearchHttp.REL_SUCCESS, 1); runner.assertTransferCount(ScrollElasticsearchHttp.REL_FAILURE, 0); } @Test public void testScrollElasticsearchOnTriggerWithServerFailNoIncomingFlowFile() throws IOException { ScrollElasticsearchHttpTestProcessor processor = new ScrollElasticsearchHttpTestProcessor(); processor.setStatus(100, "Should fail", 1); runner = TestRunners.newTestRunner(processor); // simulate doc not found runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.setValidateExpressionUsage(true); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.setIncomingConnection(false); runner.run(1, true, true); // This test generates a HTTP 100 with no incoming flow file, so nothing should be transferred processor.getRelationships().forEach(relationship -> runner.assertTransferCount(relationship, 0)); runner.assertTransferCount(ScrollElasticsearchHttp.REL_FAILURE, 0); } @Test public void testSetupSecureClient() throws Exception { ScrollElasticsearchHttpTestProcessor processor = new ScrollElasticsearchHttpTestProcessor(); runner = TestRunners.newTestRunner(processor); SSLContextService sslService = mock(SSLContextService.class); when(sslService.getIdentifier()).thenReturn("ssl-context"); runner.addControllerService("ssl-context", sslService); runner.enableControllerService(sslService); runner.setProperty(ScrollElasticsearchHttp.PROP_SSL_CONTEXT_SERVICE, "ssl-context"); runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.setValidateExpressionUsage(true); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.setIncomingConnection(false); // Allow time for the controller service to fully initialize Thread.sleep(500); runner.enqueue("".getBytes(), new HashMap<String, String>() { { put("doc_id", "28039652140"); } }); runner.run(1, true, true); } @Test public void testScrollElasticsearchOnTriggerWithIOException() throws IOException { ScrollElasticsearchHttpTestProcessor processor = new ScrollElasticsearchHttpTestProcessor(); processor.setExceptionToThrow(new IOException("Error reading from disk")); runner = TestRunners.newTestRunner(processor); // simulate doc not found runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.setValidateExpressionUsage(true); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.enqueue("".getBytes(), new HashMap<String, String>() { { put("identifier", "28039652140"); } }); runner.run(1, true, true); // This test generates a HTTP 100 "Should fail" runner.assertTransferCount(ScrollElasticsearchHttp.REL_SUCCESS, 0); runner.assertTransferCount(ScrollElasticsearchHttp.REL_FAILURE, 0); } @Test public void testScrollElasticsearchOnTriggerWithOtherException() throws IOException { ScrollElasticsearchHttpTestProcessor processor = new ScrollElasticsearchHttpTestProcessor(); processor.setExceptionToThrow(new IllegalArgumentException("Error reading from disk")); runner = TestRunners.newTestRunner(processor); // simulate doc not found runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200"); runner.setProperty(ScrollElasticsearchHttp.INDEX, "doc"); runner.setProperty(ScrollElasticsearchHttp.TYPE, "status"); runner.setValidateExpressionUsage(true); runner.setProperty(ScrollElasticsearchHttp.QUERY, "${doc_id}"); runner.enqueue("".getBytes(), new HashMap<String, String>() { { put("identifier", "28039652140"); } }); runner.run(1, true, true); // This test generates a HTTP 100 "Should fail" runner.assertTransferCount(ScrollElasticsearchHttp.REL_SUCCESS, 0); runner.assertTransferCount(ScrollElasticsearchHttp.REL_FAILURE, 1); } /** * A Test class that extends the processor in order to inject/mock behavior */ private static class ScrollElasticsearchHttpTestProcessor extends ScrollElasticsearchHttp { Exception exceptionToThrow = null; OkHttpClient client; int goodStatusCode = 200; String goodStatusMessage = "OK"; int badStatusCode; String badStatusMessage; int runNumber; List<String> pages = Arrays.asList(getDoc("scroll-page1.json"), getDoc("scroll-page2.json"), getDoc("scroll-page3.json")); public void setExceptionToThrow(Exception exceptionToThrow) { this.exceptionToThrow = exceptionToThrow; } /** * Sets the status code and message for the 1st query * * @param code * The status code to return * @param message * The status message */ void setStatus(int code, String message) { this.setStatus(code, message, 1); } /** * Sets the status code and message for the runNumber-th query * * @param code * The status code to return * @param message * The status message * @param runNumber * The run number for which to set this status */ void setStatus(int code, String message, int runNumber) { badStatusCode = code; badStatusMessage = message; this.runNumber = runNumber; } @Override protected void createElasticsearchClient(ProcessContext context) throws ProcessException { client = mock(OkHttpClient.class); OngoingStubbing<Call> stub = when(client.newCall(any(Request.class))); for (int i = 0; i < pages.size(); i++) { String page = pages.get(i); if (runNumber == i + 1) { stub = mockReturnDocument(stub, page, badStatusCode, badStatusMessage); } else { stub = mockReturnDocument(stub, page, goodStatusCode, goodStatusMessage); } } } private OngoingStubbing<Call> mockReturnDocument(OngoingStubbing<Call> stub, final String document, int statusCode, String statusMessage) { return stub.thenAnswer(new Answer<Call>() { @Override public Call answer(InvocationOnMock invocationOnMock) throws Throwable { Request realRequest = (Request) invocationOnMock.getArguments()[0]; Response mockResponse = new Response.Builder().request(realRequest).protocol(Protocol.HTTP_1_1) .code(statusCode).message(statusMessage) .body(ResponseBody.create(MediaType.parse("application/json"), document)).build(); final Call call = mock(Call.class); if (exceptionToThrow != null) { when(call.execute()).thenThrow(exceptionToThrow); } else { when(call.execute()).thenReturn(mockResponse); } return call; } }); } protected OkHttpClient getClient() { return client; } } private static String getDoc(String filename) { try { return IOUtils.toString(ScrollElasticsearchHttp.class.getClassLoader().getResourceAsStream(filename)); } catch (IOException e) { System.out.println("Error reading document " + filename); return ""; } } }