nifi-nar-bundles/nifi-mongodb-bundle/nifi-mongodb-processors/src/test/java/org/apache/nifi/processors/mongodb/RunMongoAggregationIT.java - nifi - Git at Google

 /*
  * 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.mongodb;

 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.mongodb.MongoClient;
 import com.mongodb.MongoClientURI;
 import com.mongodb.client.MongoCollection;
 import org.apache.nifi.mongodb.MongoDBClientService;
 import org.apache.nifi.mongodb.MongoDBControllerService;
 import org.apache.nifi.util.MockFlowFile;
 import org.apache.nifi.util.TestRunner;
 import org.apache.nifi.util.TestRunners;
 import org.bson.Document;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;

 import java.io.IOException;
 import java.text.SimpleDateFormat;
 import java.util.Calendar;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;

 public class RunMongoAggregationIT {

     private static final String MONGO_URI = "mongodb://localhost";
     private static final String DB_NAME   = String.format("agg_test-%s", Calendar.getInstance().getTimeInMillis());
     private static final String COLLECTION_NAME = "agg_test_data";
     private static final String AGG_ATTR = "mongo.aggregation.query";

     private TestRunner runner;
     private MongoClient mongoClient;
     private Map<String, Integer> mappings;
     private Calendar now = Calendar.getInstance();

     @BeforeEach
     public void setup() {
         runner = TestRunners.newTestRunner(RunMongoAggregation.class);
         runner.setVariable("uri", MONGO_URI);
         runner.setVariable("db", DB_NAME);
         runner.setVariable("collection", COLLECTION_NAME);
         runner.setProperty(AbstractMongoProcessor.URI, "${uri}");
         runner.setProperty(AbstractMongoProcessor.DATABASE_NAME, "${db}");
         runner.setProperty(AbstractMongoProcessor.COLLECTION_NAME, "${collection}");
         runner.setProperty(RunMongoAggregation.QUERY_ATTRIBUTE, AGG_ATTR);

         mongoClient = new MongoClient(new MongoClientURI(MONGO_URI));

         MongoCollection<Document> collection = mongoClient.getDatabase(DB_NAME).getCollection(COLLECTION_NAME);
         String[] values = new String[] { "a", "b", "c" };
         mappings = new HashMap<>();

         for (int x = 0; x < values.length; x++) {
             for (int y = 0; y < x + 2; y++) {
                 Document doc = new Document().append("val", values[x]).append("date", now.getTime());
                 collection.insertOne(doc);
             }
             mappings.put(values[x], x + 2);
         }
     }

     @AfterEach
     public void teardown() {
         runner = null;
         mongoClient.getDatabase(DB_NAME).drop();
     }

     @Test
     public void testAggregation() throws Exception {
         final String queryInput = "[\n" +
                 "    {\n" +
                 "        \"$project\": {\n" +
                 "            \"_id\": 0,\n" +
                 "            \"val\": 1\n" +
                 "        }\n" +
                 "    },\n" +
                 "    {\n" +
                 "        \"$group\": {\n" +
                 "            \"_id\": \"$val\",\n" +
                 "            \"doc_count\": {\n" +
                 "                \"$sum\": 1\n" +
                 "            }\n" +
                 "        }\n" +
                 "    }\n" +
                 "]";
         runner.setProperty(RunMongoAggregation.QUERY, queryInput);
         runner.enqueue("test");
         runner.run(1, true, true);

         evaluateRunner(1);

         runner.clearTransferState();

         runner.setIncomingConnection(false);
         runner.run(); //Null parent flowfile
         evaluateRunner(0);

         runner.run();
         List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
         for (MockFlowFile mff : flowFiles) {
             String val = mff.getAttribute(AGG_ATTR);
             Assertions.assertNotNull("Missing query attribute", val);
             Assertions.assertEquals(val, queryInput, "Value was wrong");
         }
     }

     @Test
     public void testExpressionLanguageSupport() throws Exception {
         runner.setVariable("fieldName", "$val");
         runner.setProperty(RunMongoAggregation.QUERY, "[\n" +
                 "    {\n" +
                 "        \"$project\": {\n" +
                 "            \"_id\": 0,\n" +
                 "            \"val\": 1\n" +
                 "        }\n" +
                 "    },\n" +
                 "    {\n" +
                 "        \"$group\": {\n" +
                 "            \"_id\": \"${fieldName}\",\n" +
                 "            \"doc_count\": {\n" +
                 "                \"$sum\": 1\n" +
                 "            }\n" +
                 "        }\n" +
                 "    }\n" +
                 "]");

         runner.enqueue("test");
         runner.run(1, true, true);
         evaluateRunner(1);
     }

     @Test
     public void testInvalidQuery(){
         runner.setProperty(RunMongoAggregation.QUERY, "[\n" +
             "    {\n" +
                 "        \"$invalid_stage\": {\n" +
                 "            \"_id\": 0,\n" +
                 "            \"val\": 1\n" +
                 "        }\n" +
                 "    }\n" +
             "]"
         );
         runner.enqueue("test");
         runner.run(1, true, true);
         runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, 0);
         runner.assertTransferCount(RunMongoAggregation.REL_ORIGINAL, 0);
         runner.assertTransferCount(RunMongoAggregation.REL_FAILURE, 1);
     }

     @Test
     public void testJsonTypes() throws IOException {

         runner.setProperty(RunMongoAggregation.JSON_TYPE, RunMongoAggregation.JSON_STANDARD);
         runner.setProperty(RunMongoAggregation.QUERY, "[ { \"$project\": { \"myArray\": [ \"$val\", \"$date\" ] } } ]");
         runner.enqueue("test");
         runner.run(1, true, true);

         List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
         ObjectMapper mapper = new ObjectMapper();
         SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'");
         for (MockFlowFile mockFlowFile : flowFiles) {
             byte[] raw = runner.getContentAsByteArray(mockFlowFile);
             Map<String, List<String>> read = mapper.readValue(raw, Map.class);
             Assertions.assertTrue(read.get("myArray").get(1).equalsIgnoreCase( format.format(now.getTime())));
         }

         runner.clearTransferState();

         runner.setProperty(RunMongoAggregation.JSON_TYPE, RunMongoAggregation.JSON_EXTENDED);
         runner.enqueue("test");
         runner.run(1, true, true);

         flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
         for (MockFlowFile mockFlowFile : flowFiles) {
             byte[] raw = runner.getContentAsByteArray(mockFlowFile);
             Map<String, List<Long>> read = mapper.readValue(raw, Map.class);
             Assertions.assertTrue(read.get("myArray").get(1) == now.getTimeInMillis());
         }
     }

     private void evaluateRunner(int original) throws IOException {
         runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, mappings.size());
         runner.assertTransferCount(RunMongoAggregation.REL_ORIGINAL, original);
         List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
         ObjectMapper mapper = new ObjectMapper();
         for (MockFlowFile mockFlowFile : flowFiles) {
             byte[] raw = runner.getContentAsByteArray(mockFlowFile);
             Map read = mapper.readValue(raw, Map.class);
             Assertions.assertTrue(mappings.containsKey(read.get("_id")), "Value was not found");

             String queryAttr = mockFlowFile.getAttribute(AGG_ATTR);
             Assertions.assertNotNull("Query attribute was null.", queryAttr);
             Assertions.assertTrue(queryAttr.contains("$project"), "Missing $project");
             Assertions.assertTrue(queryAttr.contains("$group"), "Missing $group");
         }
     }

     @Test
     public void testClientService() throws Exception {
         MongoDBClientService clientService = new MongoDBControllerService();
         runner.addControllerService("clientService", clientService);
         runner.removeProperty(RunMongoAggregation.URI);
         runner.setProperty(clientService, MongoDBControllerService.URI, MONGO_URI);
         runner.setProperty(RunMongoAggregation.CLIENT_SERVICE, "clientService");
         runner.setProperty(RunMongoAggregation.QUERY, "[\n" +
                         "    {\n" +
                         "        \"$project\": {\n" +
                         "            \"_id\": 0,\n" +
                         "            \"val\": 1\n" +
                         "        }\n" +
                         "    }]");
         runner.enableControllerService(clientService);
         runner.assertValid();

         runner.enqueue("{}");
         runner.run();
         runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, 9);
     }

     @Test
     public void testExtendedJsonSupport() throws Exception {
         String pattern = "yyyy-MM-dd'T'HH:mm:ss'Z'";
         SimpleDateFormat simpleDateFormat = new SimpleDateFormat(pattern);
         //Let's put this a week from now to make sure that we're not getting too close to
         //the creation date
         Date nowish = new Date(now.getTime().getTime() + (7 * 24 * 60 * 60 * 1000));

         final String queryInput = "[\n" +
             "  {\n" +
             "    \"$match\": {\n" +
             "      \"date\": { \"$gte\": { \"$date\": \"2019-01-01T00:00:00Z\" }, \"$lte\": { \"$date\": \"" + simpleDateFormat.format(nowish) + "\" } }\n" +
             "    }\n" +
             "  },\n" +
             "  {\n" +
             "    \"$group\": {\n" +
             "      \"_id\": \"$val\",\n" +
             "      \"doc_count\": {\n" +
             "        \"$sum\": 1\n" +
             "      }\n" +
             "    }\n" +
             "  }\n" +
             "]\n";

         runner.setProperty(RunMongoAggregation.QUERY, queryInput);
         runner.enqueue("test");
         runner.run(1, true, true);

         runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, mappings.size());
     }

     @Test
     public void testEmptyResponse() throws Exception {
         final String queryInput = "[\n" +
                 "  {\n" +
                 "    \"$match\": {\n" +
                 "      \"val\": \"no_exists\"\n" +
                 "    }\n" +
                 "  },\n" +
                 "  {\n" +
                 "    \"$group\": {\n" +
                 "      \"_id\": \"null\",\n" +
                 "      \"doc_count\": {\n" +
                 "        \"$sum\": 1\n" +
                 "      }\n" +
                 "    }\n" +
                 "  }\n" +
                 "]";

         runner.setProperty(RunMongoAggregation.QUERY, queryInput);
         runner.enqueue("test");
         runner.run(1, true, true);

         runner.assertTransferCount(RunMongoAggregation.REL_ORIGINAL, 1);
         runner.assertTransferCount(RunMongoAggregation.REL_FAILURE, 0);
         runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, 1);
     }
 }
	/*
	* 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.mongodb;

	import com.fasterxml.jackson.databind.ObjectMapper;
	import com.mongodb.MongoClient;
	import com.mongodb.MongoClientURI;
	import com.mongodb.client.MongoCollection;
	import org.apache.nifi.mongodb.MongoDBClientService;
	import org.apache.nifi.mongodb.MongoDBControllerService;
	import org.apache.nifi.util.MockFlowFile;
	import org.apache.nifi.util.TestRunner;
	import org.apache.nifi.util.TestRunners;
	import org.bson.Document;
	import org.junit.jupiter.api.AfterEach;
	import org.junit.jupiter.api.Assertions;
	import org.junit.jupiter.api.BeforeEach;
	import org.junit.jupiter.api.Test;

	import java.io.IOException;
	import java.text.SimpleDateFormat;
	import java.util.Calendar;
	import java.util.Date;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;

	public class RunMongoAggregationIT {

	private static final String MONGO_URI = "mongodb://localhost";
	private static final String DB_NAME = String.format("agg_test-%s", Calendar.getInstance().getTimeInMillis());
	private static final String COLLECTION_NAME = "agg_test_data";
	private static final String AGG_ATTR = "mongo.aggregation.query";

	private TestRunner runner;
	private MongoClient mongoClient;
	private Map<String, Integer> mappings;
	private Calendar now = Calendar.getInstance();

	@BeforeEach
	public void setup() {
	runner = TestRunners.newTestRunner(RunMongoAggregation.class);
	runner.setVariable("uri", MONGO_URI);
	runner.setVariable("db", DB_NAME);
	runner.setVariable("collection", COLLECTION_NAME);
	runner.setProperty(AbstractMongoProcessor.URI, "${uri}");
	runner.setProperty(AbstractMongoProcessor.DATABASE_NAME, "${db}");
	runner.setProperty(AbstractMongoProcessor.COLLECTION_NAME, "${collection}");
	runner.setProperty(RunMongoAggregation.QUERY_ATTRIBUTE, AGG_ATTR);

	mongoClient = new MongoClient(new MongoClientURI(MONGO_URI));

	MongoCollection<Document> collection = mongoClient.getDatabase(DB_NAME).getCollection(COLLECTION_NAME);
	String[] values = new String[] { "a", "b", "c" };
	mappings = new HashMap<>();

	for (int x = 0; x < values.length; x++) {
	for (int y = 0; y < x + 2; y++) {
	Document doc = new Document().append("val", values[x]).append("date", now.getTime());
	collection.insertOne(doc);
	}
	mappings.put(values[x], x + 2);
	}
	}

	@AfterEach
	public void teardown() {
	runner = null;
	mongoClient.getDatabase(DB_NAME).drop();
	}

	@Test
	public void testAggregation() throws Exception {
	final String queryInput = "[\n" +
	" {\n" +
	" \"$project\": {\n" +
	" \"_id\": 0,\n" +
	" \"val\": 1\n" +
	" }\n" +
	" },\n" +
	" {\n" +
	" \"$group\": {\n" +
	" \"_id\": \"$val\",\n" +
	" \"doc_count\": {\n" +
	" \"$sum\": 1\n" +
	" }\n" +
	" }\n" +
	" }\n" +
	"]";
	runner.setProperty(RunMongoAggregation.QUERY, queryInput);
	runner.enqueue("test");
	runner.run(1, true, true);

	evaluateRunner(1);

	runner.clearTransferState();

	runner.setIncomingConnection(false);
	runner.run(); //Null parent flowfile
	evaluateRunner(0);

	runner.run();
	List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
	for (MockFlowFile mff : flowFiles) {
	String val = mff.getAttribute(AGG_ATTR);
	Assertions.assertNotNull("Missing query attribute", val);
	Assertions.assertEquals(val, queryInput, "Value was wrong");
	}
	}

	@Test
	public void testExpressionLanguageSupport() throws Exception {
	runner.setVariable("fieldName", "$val");
	runner.setProperty(RunMongoAggregation.QUERY, "[\n" +
	" {\n" +
	" \"$project\": {\n" +
	" \"_id\": 0,\n" +
	" \"val\": 1\n" +
	" }\n" +
	" },\n" +
	" {\n" +
	" \"$group\": {\n" +
	" \"_id\": \"${fieldName}\",\n" +
	" \"doc_count\": {\n" +
	" \"$sum\": 1\n" +
	" }\n" +
	" }\n" +
	" }\n" +
	"]");

	runner.enqueue("test");
	runner.run(1, true, true);
	evaluateRunner(1);
	}

	@Test
	public void testInvalidQuery(){
	runner.setProperty(RunMongoAggregation.QUERY, "[\n" +
	" {\n" +
	" \"$invalid_stage\": {\n" +
	" \"_id\": 0,\n" +
	" \"val\": 1\n" +
	" }\n" +
	" }\n" +
	"]"
	);
	runner.enqueue("test");
	runner.run(1, true, true);
	runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, 0);
	runner.assertTransferCount(RunMongoAggregation.REL_ORIGINAL, 0);
	runner.assertTransferCount(RunMongoAggregation.REL_FAILURE, 1);
	}

	@Test
	public void testJsonTypes() throws IOException {

	runner.setProperty(RunMongoAggregation.JSON_TYPE, RunMongoAggregation.JSON_STANDARD);
	runner.setProperty(RunMongoAggregation.QUERY, "[ { \"$project\": { \"myArray\": [ \"$val\", \"$date\" ] } } ]");
	runner.enqueue("test");
	runner.run(1, true, true);

	List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
	ObjectMapper mapper = new ObjectMapper();
	SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'");
	for (MockFlowFile mockFlowFile : flowFiles) {
	byte[] raw = runner.getContentAsByteArray(mockFlowFile);
	Map<String, List<String>> read = mapper.readValue(raw, Map.class);
	Assertions.assertTrue(read.get("myArray").get(1).equalsIgnoreCase( format.format(now.getTime())));
	}

	runner.clearTransferState();

	runner.setProperty(RunMongoAggregation.JSON_TYPE, RunMongoAggregation.JSON_EXTENDED);
	runner.enqueue("test");
	runner.run(1, true, true);

	flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
	for (MockFlowFile mockFlowFile : flowFiles) {
	byte[] raw = runner.getContentAsByteArray(mockFlowFile);
	Map<String, List<Long>> read = mapper.readValue(raw, Map.class);
	Assertions.assertTrue(read.get("myArray").get(1) == now.getTimeInMillis());
	}
	}

	private void evaluateRunner(int original) throws IOException {
	runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, mappings.size());
	runner.assertTransferCount(RunMongoAggregation.REL_ORIGINAL, original);
	List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(RunMongoAggregation.REL_RESULTS);
	ObjectMapper mapper = new ObjectMapper();
	for (MockFlowFile mockFlowFile : flowFiles) {
	byte[] raw = runner.getContentAsByteArray(mockFlowFile);
	Map read = mapper.readValue(raw, Map.class);
	Assertions.assertTrue(mappings.containsKey(read.get("_id")), "Value was not found");

	String queryAttr = mockFlowFile.getAttribute(AGG_ATTR);
	Assertions.assertNotNull("Query attribute was null.", queryAttr);
	Assertions.assertTrue(queryAttr.contains("$project"), "Missing $project");
	Assertions.assertTrue(queryAttr.contains("$group"), "Missing $group");
	}
	}

	@Test
	public void testClientService() throws Exception {
	MongoDBClientService clientService = new MongoDBControllerService();
	runner.addControllerService("clientService", clientService);
	runner.removeProperty(RunMongoAggregation.URI);
	runner.setProperty(clientService, MongoDBControllerService.URI, MONGO_URI);
	runner.setProperty(RunMongoAggregation.CLIENT_SERVICE, "clientService");
	runner.setProperty(RunMongoAggregation.QUERY, "[\n" +
	" {\n" +
	" \"$project\": {\n" +
	" \"_id\": 0,\n" +
	" \"val\": 1\n" +
	" }\n" +
	" }]");
	runner.enableControllerService(clientService);
	runner.assertValid();

	runner.enqueue("{}");
	runner.run();
	runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, 9);
	}

	@Test
	public void testExtendedJsonSupport() throws Exception {
	String pattern = "yyyy-MM-dd'T'HH:mm:ss'Z'";
	SimpleDateFormat simpleDateFormat = new SimpleDateFormat(pattern);
	//Let's put this a week from now to make sure that we're not getting too close to
	//the creation date
	Date nowish = new Date(now.getTime().getTime() + (7 * 24 * 60 * 60 * 1000));

	final String queryInput = "[\n" +
	" {\n" +
	" \"$match\": {\n" +
	" \"date\": { \"$gte\": { \"$date\": \"2019-01-01T00:00:00Z\" }, \"$lte\": { \"$date\": \"" + simpleDateFormat.format(nowish) + "\" } }\n" +
	" }\n" +
	" },\n" +
	" {\n" +
	" \"$group\": {\n" +
	" \"_id\": \"$val\",\n" +
	" \"doc_count\": {\n" +
	" \"$sum\": 1\n" +
	" }\n" +
	" }\n" +
	" }\n" +
	"]\n";

	runner.setProperty(RunMongoAggregation.QUERY, queryInput);
	runner.enqueue("test");
	runner.run(1, true, true);

	runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, mappings.size());
	}

	@Test
	public void testEmptyResponse() throws Exception {
	final String queryInput = "[\n" +
	" {\n" +
	" \"$match\": {\n" +
	" \"val\": \"no_exists\"\n" +
	" }\n" +
	" },\n" +
	" {\n" +
	" \"$group\": {\n" +
	" \"_id\": \"null\",\n" +
	" \"doc_count\": {\n" +
	" \"$sum\": 1\n" +
	" }\n" +
	" }\n" +
	" }\n" +
	"]";

	runner.setProperty(RunMongoAggregation.QUERY, queryInput);
	runner.enqueue("test");
	runner.run(1, true, true);

	runner.assertTransferCount(RunMongoAggregation.REL_ORIGINAL, 1);
	runner.assertTransferCount(RunMongoAggregation.REL_FAILURE, 0);
	runner.assertTransferCount(RunMongoAggregation.REL_RESULTS, 1);
	}
	}