Java Code Examples for org.apache.nifi.util.TestRunner#enableControllerService()

@Test(timeout = 10000)
public void whenExceptionIsRaisedTheProcessorShouldBeYielded() throws Exception {
    TestRunner runner = TestRunners.newTestRunner(new ConsumeJMS());
    JMSConnectionFactoryProviderDefinition cs = mock(JMSConnectionFactoryProviderDefinition.class);
    ActiveMQConnectionFactory cf = new ActiveMQConnectionFactory("tcp://invalidhost:9999?soTimeout=3");

    when(cs.getIdentifier()).thenReturn("cfProvider");
    when(cs.getConnectionFactory()).thenReturn(cf);
    runner.addControllerService("cfProvider", cs);
    runner.enableControllerService(cs);

    runner.setProperty(ConsumeJMS.CF_SERVICE, "cfProvider");
    runner.setProperty(ConsumeJMS.DESTINATION, "foo");
    runner.setProperty(ConsumeJMS.DESTINATION_TYPE, ConsumeJMS.TOPIC);

    try {
        runner.run();
        fail("The test was implemented in a way this line should not be reached.");
    } catch (AssertionError e) {
    } finally {
        assertTrue("In case of an exception, the processor should be yielded.", ((MockProcessContext) runner.getProcessContext()).isYieldCalled());
    }
}
 

@Test
public void testJsonSpelRules() throws InitializationException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);
    final RulesEngineService service = new MockEasyRulesEngineService();
    runner.addControllerService("easy-rules-engine-service-test",service);
    runner.setProperty(service, EasyRulesEngineService.RULES_FILE_PATH, "src/test/resources/test_spel_rules.json");
    runner.setProperty(service,EasyRulesEngineService.RULES_FILE_TYPE, "JSON");
    runner.setProperty(service,EasyRulesEngineService.RULES_FILE_FORMAT, "SPEL");
    runner.enableControllerService(service);
    runner.assertValid(service);
    Map<String, Object> facts = new HashMap<>();
    facts.put("predictedQueuedCount",60);
    facts.put("predictedTimeToBytesBackpressureMillis",299999);
    List<Action> actions = service.fireRules(facts);
    assertNotNull(actions);
    assertEquals(actions.size(), 2);
}
 

@Test
public void testWithExceptionThrown() throws InitializationException {
    final TestRunner runner = TestRunners.newTestRunner(HandleHttpResponse.class);

    final MockHttpContextMap contextMap = new MockHttpContextMap("my-id", "FlowFileAccessException");
    runner.addControllerService("http-context-map", contextMap);
    runner.enableControllerService(contextMap);
    runner.setProperty(HandleHttpResponse.HTTP_CONTEXT_MAP, "http-context-map");
    runner.setProperty(HandleHttpResponse.STATUS_CODE, "${status.code}");
    runner.setProperty("my-attr", "${my-attr}");
    runner.setProperty("no-valid-attr", "${no-valid-attr}");

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(HTTPUtils.HTTP_CONTEXT_ID, "my-id");
    attributes.put("my-attr", "hello");
    attributes.put("status.code", "201");

    runner.enqueue("hello".getBytes(), attributes);

    runner.run();

    runner.assertAllFlowFilesTransferred(HandleHttpResponse.REL_FAILURE, 1);
    assertEquals(0, contextMap.getCompletionCount());
}
 

@Test
public void testKeysCredentialsProviderWithRoleAndName() throws Throwable {
    final TestRunner runner = TestRunners.newTestRunner(FetchS3Object.class);
    final AWSCredentialsProviderControllerService serviceImpl = new AWSCredentialsProviderControllerService();
    runner.addControllerService("awsCredentialsProvider", serviceImpl);
    runner.setProperty(serviceImpl, AbstractAWSProcessor.ACCESS_KEY, "awsAccessKey");
    runner.setProperty(serviceImpl, AbstractAWSProcessor.SECRET_KEY, "awsSecretKey");
    runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.ASSUME_ROLE_ARN, "Role");
    runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.ASSUME_ROLE_NAME, "RoleName");
    runner.enableControllerService(serviceImpl);

    runner.assertValid(serviceImpl);
    final AWSCredentialsProviderService service = (AWSCredentialsProviderService) runner.getProcessContext()
            .getControllerServiceLookup().getControllerService("awsCredentialsProvider");
    Assert.assertNotNull(service);
    final AWSCredentialsProvider credentialsProvider = service.getCredentialsProvider();
    Assert.assertNotNull(credentialsProvider);
    assertEquals("credentials provider should be equal", STSAssumeRoleSessionCredentialsProvider.class,
            credentialsProvider.getClass());
}
 

@Test
public void testFailInMiddleWithBadParameterType() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutSQL.class);
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    runner.setProperty(PutSQL.OBTAIN_GENERATED_KEYS, "false");
    runner.setProperty(PutSQL.CONNECTION_POOL, "dbcp");

    final Map<String, String> goodAttributes = new HashMap<>();
    goodAttributes.put("sql.args.1.type", String.valueOf(Types.INTEGER));
    goodAttributes.put("sql.args.1.value", "84");

    final Map<String, String> badAttributes = new HashMap<>();
    badAttributes.put("sql.args.1.type", String.valueOf(Types.VARCHAR));
    badAttributes.put("sql.args.1.value", "hello");

    final byte[] data = "INSERT INTO PERSONS_AI (NAME, CODE) VALUES ('Mark', ?)".getBytes();
    runner.enqueue(data, goodAttributes);
    runner.enqueue(data, badAttributes);
    runner.enqueue(data, goodAttributes);
    runner.enqueue(data, goodAttributes);
    runner.run();

    runner.assertTransferCount(PutSQL.REL_FAILURE, 1);
    runner.assertTransferCount(PutSQL.REL_SUCCESS, 3);
}
 

@Test
public void testYamlNiFiRules() throws InitializationException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);
    final RulesEngineService service = new MockEasyRulesEngineService();
    runner.addControllerService("easy-rules-engine-service-test",service);
    runner.setProperty(service, EasyRulesEngineService.RULES_FILE_PATH, "src/test/resources/test_nifi_rules.yml");
    runner.setProperty(service,EasyRulesEngineService.RULES_FILE_TYPE, "YAML");
    runner.setProperty(service,EasyRulesEngineService.RULES_FILE_FORMAT, "NIFI");
    runner.enableControllerService(service);
    runner.assertValid(service);
    Map<String, Object> facts = new HashMap<>();
    facts.put("predictedQueuedCount",60);
    facts.put("predictedTimeToBytesBackpressureMillis",299999);
    List<Action> actions = service.fireRules(facts);
    assertNotNull(actions);
    assertEquals(actions.size(), 3);
}
 

@Test
public void testFailAtBeginning() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHive3QL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersonsAutoId);
        }
    }

    runner.setProperty(PutHive3QL.HIVE_DBCP_SERVICE, "dbcp");
    runner.enqueue("INSERT INTO PERSONS".getBytes()); // intentionally wrong syntax
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Tom', 3)".getBytes());
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Harry', 44)".getBytes());
    runner.run();

    runner.assertTransferCount(PutHive3QL.REL_FAILURE, 1);
    runner.assertTransferCount(PutHive3QL.REL_SUCCESS, 2);
}
 

private AtomicDistributedMapCacheClient<byte[]> configureHBaseCacheService(final TestRunner runner, final HBaseClientService service) throws InitializationException {
    final HBase_1_1_2_ClientMapCacheService cacheService = new HBase_1_1_2_ClientMapCacheService();
    runner.addControllerService("hbaseCache", cacheService);
    runner.setProperty(cacheService, HBase_1_1_2_ClientMapCacheService.HBASE_CLIENT_SERVICE, "hbaseClient");
    runner.setProperty(cacheService, HBase_1_1_2_ClientMapCacheService.HBASE_CACHE_TABLE_NAME, tableName);
    runner.setProperty(cacheService, HBase_1_1_2_ClientMapCacheService.HBASE_COLUMN_FAMILY, columnFamily);
    runner.setProperty(cacheService, HBase_1_1_2_ClientMapCacheService.HBASE_COLUMN_QUALIFIER, columnQualifier);
    runner.enableControllerService(cacheService);
    runner.setProperty(TestProcessor.HBASE_CACHE_SERVICE,"hbaseCache");
    return cacheService;
}
 

@Test
public void testRetryableFailure() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHiveQL.class);
    final DBCPService service = new SQLExceptionService(null);
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);

    runner.setProperty(PutHiveQL.HIVE_DBCP_SERVICE, "dbcp");

    final String sql = "INSERT INTO PERSONS (ID, NAME, CODE) VALUES (?, ?, ?); " +
        "UPDATE PERSONS SET NAME='George' WHERE ID=?; ";

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("hiveql.args.1.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.1.value", "1");

    attributes.put("hiveql.args.2.type", String.valueOf(Types.VARCHAR));
    attributes.put("hiveql.args.2.value", "Mark");

    attributes.put("hiveql.args.3.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.3.value", "84");

    attributes.put("hiveql.args.4.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.4.value", "1");

    runner.enqueue(sql.getBytes(), attributes);
    runner.run();

    // should fail because there isn't a valid connection and tables don't exist.
    runner.assertAllFlowFilesTransferred(PutHiveQL.REL_RETRY, 1);
}
 

@Test
public void testTransformControllerNoKey() throws IOException, InitializationException {
    final TestRunner runner = TestRunners.newTestRunner(new TransformXml());

    final SimpleKeyValueLookupService service = new SimpleKeyValueLookupService();
    runner.addControllerService("simple-key-value-lookup-service", service);
    runner.setProperty(service, "key1", "value1");
    runner.enableControllerService(service);
    runner.assertValid(service);
    runner.setProperty(TransformXml.XSLT_CONTROLLER, "simple-key-value-lookup-service");

    runner.assertNotValid();
}
 

@Test
public void testTransformCalc() throws InitializationException, IOException, SQLException {
    final MockRecordParser parser = new MockRecordParser();
    parser.addSchemaField("ID", RecordFieldType.INT);
    parser.addSchemaField("AMOUNT1", RecordFieldType.FLOAT);
    parser.addSchemaField("AMOUNT2", RecordFieldType.FLOAT);
    parser.addSchemaField("AMOUNT3", RecordFieldType.FLOAT);

    parser.addRecord(8, 10.05F, 15.45F, 89.99F);
    parser.addRecord(100, 20.25F, 25.25F, 45.25F);
    parser.addRecord(105, 20.05F, 25.05F, 45.05F);
    parser.addRecord(200, 34.05F, 25.05F, 75.05F);

    final MockRecordWriter writer = new MockRecordWriter("\"NAME\",\"POINTS\"");

    TestRunner runner = getRunner();
    runner.addControllerService("parser", parser);
    runner.enableControllerService(parser);
    runner.addControllerService("writer", writer);
    runner.enableControllerService(writer);

    runner.setProperty(REL_NAME, "select ID, AMOUNT1+AMOUNT2+AMOUNT3 as TOTAL from FLOWFILE where ID=100");
    runner.setProperty(QueryRecord.RECORD_READER_FACTORY, "parser");
    runner.setProperty(QueryRecord.RECORD_WRITER_FACTORY, "writer");

    runner.enqueue(new byte[0]);
    runner.run();

    runner.assertTransferCount(REL_NAME, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(REL_NAME).get(0);

    out.assertContentEquals("\"NAME\",\"POINTS\"\n\"100\",\"90.75\"\n");
}
 

@Test
public void testUpdateWithMissingColumnWarning() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(ConvertJSONToSQL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersons);
        }
    }

    runner.setProperty(ConvertJSONToSQL.CONNECTION_POOL, "dbcp");
    runner.setProperty(ConvertJSONToSQL.TABLE_NAME, "PERSONS");
    runner.setProperty(ConvertJSONToSQL.STATEMENT_TYPE, "UPDATE");
    runner.setProperty(ConvertJSONToSQL.UPDATE_KEY, "name,  code, extra");
    runner.setProperty(ConvertJSONToSQL.UNMATCHED_COLUMN_BEHAVIOR, ConvertJSONToSQL.WARNING_UNMATCHED_COLUMN);
    runner.enqueue(Paths.get("src/test/resources/TestConvertJSONToSQL/person-1.json"));
    runner.run();

    runner.assertTransferCount(ConvertJSONToSQL.REL_ORIGINAL, 1);
    runner.getFlowFilesForRelationship(ConvertJSONToSQL.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT.key(), "1");
    runner.assertTransferCount(ConvertJSONToSQL.REL_SQL, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ConvertJSONToSQL.REL_SQL).get(0);
    out.assertAttributeEquals("sql.args.1.type", String.valueOf(java.sql.Types.INTEGER));
    out.assertAttributeEquals("sql.args.1.value", "1");
    out.assertAttributeEquals("sql.args.2.type", String.valueOf(java.sql.Types.VARCHAR));
    out.assertAttributeEquals("sql.args.2.value", "Mark");
    out.assertAttributeEquals("sql.args.3.type", String.valueOf(java.sql.Types.INTEGER));
    out.assertAttributeEquals("sql.args.3.value", "48");

    out.assertContentEquals("UPDATE PERSONS SET ID = ? WHERE NAME = ? AND CODE = ?");

}
 

@Test
public void testMultipleStatementsWithinFlowFile() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutSQL.class);
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    runner.setProperty(PutSQL.CONNECTION_POOL, "dbcp");

    recreateTable("PERSONS", createPersons);

    final String sql = "INSERT INTO PERSONS (ID, NAME, CODE) VALUES (?, ?, ?); " +
            "UPDATE PERSONS SET NAME='George' WHERE ID=?; ";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("sql.args.1.type", String.valueOf(Types.INTEGER));
    attributes.put("sql.args.1.value", "1");

    attributes.put("sql.args.2.type", String.valueOf(Types.VARCHAR));
    attributes.put("sql.args.2.value", "Mark");

    attributes.put("sql.args.3.type", String.valueOf(Types.INTEGER));
    attributes.put("sql.args.3.value", "84");

    attributes.put("sql.args.4.type", String.valueOf(Types.INTEGER));
    attributes.put("sql.args.4.value", "1");

    runner.enqueue(sql.getBytes(), attributes);
    runner.run();

    // should fail because of the semicolon
    runner.assertAllFlowFilesTransferred(PutSQL.REL_FAILURE, 1);

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            final ResultSet rs = stmt.executeQuery("SELECT * FROM PERSONS");
            assertFalse(rs.next());
        }
    }
}
 

@Test
public void testFailInMiddleWithBadStatement() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHive3QL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersonsAutoId);
        }
    }

    runner.setProperty(PutHive3QL.HIVE_DBCP_SERVICE, "dbcp");
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Mark', 84)".getBytes());
    runner.enqueue("INSERT INTO PERSONS".getBytes()); // intentionally wrong syntax
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Tom', 3)".getBytes());
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Harry', 44)".getBytes());
    runner.run();

    runner.assertTransferCount(PutHive3QL.REL_FAILURE, 1);
    runner.assertTransferCount(PutHive3QL.REL_SUCCESS, 3);
    runner.getFlowFilesForRelationship(PutHive3QL.REL_SUCCESS)
            .forEach(f -> f.assertAttributeEquals(PutHive3QL.ATTR_OUTPUT_TABLES, "PERSONS"));
}
 

private MockHBaseClientService getHBaseClientService(TestRunner runner) throws InitializationException {
    final MockHBaseClientService hBaseClient = new MockHBaseClientService();
    runner.addControllerService("hbaseClient", hBaseClient);
    runner.enableControllerService(hBaseClient);
    runner.setProperty(PutHBaseCell.HBASE_CLIENT_SERVICE, "hbaseClient");
    return hBaseClient;
}
 

@Test
public void testKeysCredentialsProviderWithRoleAndNameAndSessionTimeout3600() throws Throwable {
    final TestRunner runner = TestRunners.newTestRunner(FetchS3Object.class);
    final AWSCredentialsProviderControllerService serviceImpl = new AWSCredentialsProviderControllerService();
    runner.addControllerService("awsCredentialsProvider", serviceImpl);
    runner.setProperty(serviceImpl, AbstractAWSProcessor.ACCESS_KEY, "awsAccessKey");
    runner.setProperty(serviceImpl, AbstractAWSProcessor.SECRET_KEY, "awsSecretKey");
    runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.ASSUME_ROLE_ARN, "Role");
    runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.ASSUME_ROLE_NAME, "RoleName");
    runner.setProperty(serviceImpl, AWSCredentialsProviderControllerService.MAX_SESSION_TIME, "900");
    runner.enableControllerService(serviceImpl);

    runner.assertValid(serviceImpl);
}
 

@Test
public void testService() throws InitializationException {
    final TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);
    final StandardMyService service = new StandardMyService();
    runner.addControllerService("test-good", service);

    runner.setProperty(service, StandardMyService.MY_PROPERTY, "test-value");
    runner.enableControllerService(service);

    runner.assertValid(service);
}
 

@Test
public void testInvalidStatement() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHive_1_1QL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersons);
        }
    }

    runner.setProperty(PutHive_1_1QL.HIVE_DBCP_SERVICE, "dbcp");

    final String sql = "INSERT INTO PERSONS (ID, NAME, CODE) VALUES (?, ?, ?); " +
        "UPDATE SOME_RANDOM_TABLE NAME='George' WHERE ID=?; ";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("hiveql.args.1.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.1.value", "1");

    attributes.put("hiveql.args.2.type", String.valueOf(Types.VARCHAR));
    attributes.put("hiveql.args.2.value", "Mark");

    attributes.put("hiveql.args.3.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.3.value", "84");

    attributes.put("hiveql.args.4.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.4.value", "1");

    runner.enqueue(sql.getBytes(), attributes);
    runner.run();

    // should fail because of the table is invalid
    runner.assertAllFlowFilesTransferred(PutHive_1_1QL.REL_FAILURE, 1);

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            final ResultSet rs = stmt.executeQuery("SELECT * FROM PERSONS");
            assertTrue(rs.next());
        }
    }
}
 

@Test
public void testPersistentMapServerAndClientWithLFUEvictions() throws InitializationException, IOException {
    /**
     * This bypasses the test for build environments in OS X running Java 1.8 due to a JVM bug
     * See:  https://issues.apache.org/jira/browse/NIFI-437
     */
    Assume.assumeFalse("test is skipped due to build environment being OS X with JDK 1.8. See https://issues.apache.org/jira/browse/NIFI-437",
            SystemUtils.IS_OS_MAC && SystemUtils.IS_JAVA_1_8);

    LOGGER.info("Testing " + Thread.currentThread().getStackTrace()[1].getMethodName());
    // Create server
    final File dataFile = new File("target/cache-data");
    deleteRecursively(dataFile);

    // Create server
    final TestRunner runner = TestRunners.newTestRunner(Mockito.mock(Processor.class));
    final DistributedMapCacheServer server = new MapServer();
    runner.addControllerService("server", server);
    runner.setProperty(server, DistributedMapCacheServer.PERSISTENCE_PATH, dataFile.getAbsolutePath());
    runner.setProperty(server, DistributedMapCacheServer.MAX_CACHE_ENTRIES, "3");
    runner.setProperty(server, DistributedMapCacheServer.EVICTION_POLICY, DistributedMapCacheServer.EVICTION_STRATEGY_LFU);
    runner.enableControllerService(server);

    DistributedMapCacheClientService client = createMapClient(server.getPort());
    final Serializer<String> serializer = new StringSerializer();
    final boolean added = client.putIfAbsent("test", "1", serializer, serializer);
    waitABit();
    final boolean added2 = client.putIfAbsent("test2", "2", serializer, serializer);
    waitABit();
    final boolean added3 = client.putIfAbsent("test3", "3", serializer, serializer);
    waitABit();
    assertTrue(added);
    assertTrue(added2);
    assertTrue(added3);

    final boolean contains = client.containsKey("test", serializer);
    final boolean contains2 = client.containsKey("test2", serializer);
    assertTrue(contains);
    assertTrue(contains2);

    final Deserializer<String> deserializer = new StringDeserializer();
    final Set<String> keys = client.keySet(deserializer);
    assertEquals(3, keys.size());
    assertTrue(keys.contains("test"));
    assertTrue(keys.contains("test2"));
    assertTrue(keys.contains("test3"));

    final boolean addedAgain = client.putIfAbsent("test", "1", serializer, serializer);
    assertFalse(addedAgain);

    final boolean added4 = client.putIfAbsent("test4", "4", serializer, serializer);
    assertTrue(added4);

    // ensure that added3 was evicted because it was used least frequently
    assertFalse(client.containsKey("test3", serializer));

    server.shutdownServer();

    final DistributedMapCacheServer newServer = new MapServer();
    runner.addControllerService("server2", newServer);
    runner.setProperty(newServer, DistributedMapCacheServer.PERSISTENCE_PATH, dataFile.getAbsolutePath());
    runner.enableControllerService(newServer);
    client.close();
    client = createMapClient(newServer.getPort());

    assertTrue(client.containsKey("test", serializer));
    assertTrue(client.containsKey("test2", serializer));
    assertFalse(client.containsKey("test3", serializer));
    assertTrue(client.containsKey("test4", serializer));

    // Test removeByPattern, the first two should be removed and the last should remain
    client.put("test.1", "1", serializer, serializer);
    client.put("test.2", "2", serializer, serializer);
    client.put("test3", "2", serializer, serializer);
    final long removedTwo = client.removeByPattern("test\\..*");
    assertEquals(2L, removedTwo);
    assertFalse(client.containsKey("test.1", serializer));
    assertFalse(client.containsKey("test.2", serializer));
    assertTrue(client.containsKey("test3", serializer));

    // test removeByPatternAndGet
    client.put("test.1", "1", serializer, serializer);
    client.put("test.2", "2", serializer, serializer);
    Map<String,String> removed = client.removeByPatternAndGet("test\\..*", deserializer, deserializer);
    assertEquals(2, removed.size());
    assertTrue(removed.containsKey("test.1"));
    assertTrue(removed.containsKey("test.2"));
    assertFalse(client.containsKey("test.1", serializer));
    assertFalse(client.containsKey("test.2", serializer));
    assertTrue(client.containsKey("test3", serializer));
    removed = client.removeByPatternAndGet("test\\..*", deserializer, deserializer);
    assertEquals(0, removed.size());

    newServer.shutdownServer();
    client.close();
}
 

@Test
public void testSplitMode() throws InitializationException, IOException {
    String expectedOutput = null;
    final TestRunner runner = TestRunners.newTestRunner(new ForkRecord());
    final JsonTreeReader jsonReader = new JsonTreeReader();
    runner.addControllerService("record-reader", jsonReader);

    final String inputSchemaText = new String(Files.readAllBytes(Paths.get("src/test/resources/TestForkRecord/schema/schema.avsc")));
    final String outputSchemaText = new String(Files.readAllBytes(Paths.get("src/test/resources/TestForkRecord/schema/schema.avsc")));

    runner.setProperty(jsonReader, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaAccessUtils.SCHEMA_TEXT_PROPERTY);
    runner.setProperty(jsonReader, SchemaAccessUtils.SCHEMA_TEXT, inputSchemaText);
    runner.enableControllerService(jsonReader);
    runner.setProperty(ForkRecord.RECORD_READER, "record-reader");

    final JsonRecordSetWriter jsonWriter = new JsonRecordSetWriter();
    runner.addControllerService("record-writer", jsonWriter);
    runner.setProperty(jsonWriter, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaAccessUtils.SCHEMA_TEXT_PROPERTY);
    runner.setProperty(jsonWriter, SchemaAccessUtils.SCHEMA_TEXT, outputSchemaText);
    runner.setProperty(jsonWriter, "Pretty Print JSON", "true");
    runner.setProperty(jsonWriter, "Schema Write Strategy", "full-schema-attribute");
    runner.enableControllerService(jsonWriter);
    runner.setProperty(ForkRecord.RECORD_WRITER, "record-writer");

    runner.setProperty(ForkRecord.MODE, ForkRecord.MODE_SPLIT);

    runner.setProperty("my-path", "/address");
    runner.enqueue(Paths.get("src/test/resources/TestForkRecord/input/complex-input-json.json"));
    runner.run();
    runner.assertTransferCount(ForkRecord.REL_ORIGINAL, 1);
    runner.assertTransferCount(ForkRecord.REL_FORK, 1);
    expectedOutput = new String(Files.readAllBytes(Paths.get("src/test/resources/TestForkRecord/output/split-address.json")));
    runner.getFlowFilesForRelationship(ForkRecord.REL_FORK).get(0).assertContentEquals(expectedOutput);
    runner.getFlowFilesForRelationship(ForkRecord.REL_FORK).get(0).assertAttributeEquals("record.count", "5");

    runner.clearTransferState();
    runner.setProperty("my-path", "/bankAccounts[*]/last5Transactions");
    runner.enqueue(Paths.get("src/test/resources/TestForkRecord/input/complex-input-json.json"));
    runner.run();
    runner.assertTransferCount(ForkRecord.REL_ORIGINAL, 1);
    runner.assertTransferCount(ForkRecord.REL_FORK, 1);
    expectedOutput = new String(Files.readAllBytes(Paths.get("src/test/resources/TestForkRecord/output/split-transactions.json")));
    runner.getFlowFilesForRelationship(ForkRecord.REL_FORK).get(0).assertContentEquals(expectedOutput);
    runner.getFlowFilesForRelationship(ForkRecord.REL_FORK).get(0).assertAttributeEquals("record.count", "6");
}