FlowEnricher { public void enrichFlowWithBundleInformation() throws FlowEnrichmentException { final Path flowPath = niFiProperties.getFlowConfigurationFile().toPath(); logger.debug("Enriching generated {} with bundling information", flowPath.toAbsolutePath()); try { final Document flowDocument = flowParser.parse(flowPath.toAbsolutePath().toFile()); if (flowDocument == null) { throw new FlowEnrichmentException("Unable to successfully parse the specified flow at " + flowPath.toAbsolutePath()); } final Map<String, EnrichingElementAdapter> componentEnrichingMap = new HashMap<>(); for (String typeElementName : Arrays.asList(PROCESSOR_TAG_NAME, CONTROLLER_SERVICE_TAG_NAME, REPORTING_TASK_TAG_NAME)) { final NodeList componentNodeList = flowDocument.getElementsByTagName(typeElementName); final Map<String, EnrichingElementAdapter> elementIdToMetadataMap = mapComponents(componentNodeList); componentEnrichingMap.putAll(elementIdToMetadataMap); } for (Map.Entry<String, EnrichingElementAdapter> componentIdToMetadata : componentEnrichingMap.entrySet()) { final EnrichingElementAdapter componentToEnrich = componentIdToMetadata.getValue(); if (componentToEnrich.getBundleElement() != null) { continue; } final String componentToEnrichClass = componentToEnrich.getComponentClass(); final Map<String, Bundle> componentToEnrichVersionToBundles = minifi.getBundles(componentToEnrichClass) .stream() .collect(Collectors.toMap(bundle -> bundle.getBundleDetails().getCoordinate().getVersion(), bundle -> bundle)); enrichComponent(componentToEnrich, componentToEnrichVersionToBundles); } flowParser.writeFlow(flowDocument, flowPath.toAbsolutePath()); } catch (IOException | TransformerException e) { throw new FlowEnrichmentException("Unable to successfully automate the enrichment of the generated flow with bundle information", e); } } FlowEnricher(MiNiFi minifi, FlowParser flowParser, NiFiProperties niFiProperties); void enrichFlowWithBundleInformation(); static final String PROCESSOR_TAG_NAME; static final String CONTROLLER_SERVICE_TAG_NAME; static final String REPORTING_TASK_TAG_NAME; }
@Test public void enrichFlowWithBundleInformationTest() throws Exception { final FlowParser parser = new FlowParser(); final Document flowDocument = parser.parse(TEST_FLOW_XML_FILE); final FlowParser mockFlowParser = Mockito.mock(FlowParser.class); Mockito.when(mockFlowParser.parse(any())).thenReturn(flowDocument); final MiNiFi minifi = Mockito.mock(MiNiFi.class); final NiFiProperties mockProperties = Mockito.mock(NiFiProperties.class); Mockito.when(mockProperties.getFlowConfigurationFile()).thenReturn(new File(TEST_FLOW_XML_PATH)); final String defaultGroup = "org.apache.nifi"; final String aVersion = "1.0.0"; final String anotherVersion = "2.0.0"; final String minifiGroup = "org.apache.nifi.minifi"; final String minifiStandardNarId = "minifi-standard-nar"; final String minifiVersion = "0.0.2"; final String standardSvcsId = "nifi-standard-services-api-nar"; final List<Bundle> standardSvcsBundles = new ArrayList<>(); standardSvcsBundles.add(generateBundle(defaultGroup, standardSvcsId, aVersion)); standardSvcsBundles.add(generateBundle(defaultGroup, standardSvcsId, anotherVersion)); final String sslContextSvcId = "nifi-ssl-context-service-nar"; final List<Bundle> sslBundles = new ArrayList<>(); sslBundles.add(generateBundle(defaultGroup, sslContextSvcId, aVersion, standardSvcsBundles.get(0).getBundleDetails().getCoordinate())); sslBundles.add(generateBundle(defaultGroup, sslContextSvcId, anotherVersion, standardSvcsBundles.get(1).getBundleDetails().getCoordinate())); List<Bundle> minifiStdBundles = new ArrayList<>(); minifiStdBundles.add(generateBundle(minifiGroup, minifiStandardNarId, minifiVersion, standardSvcsBundles.get(0).getBundleDetails().getCoordinate())); List<Bundle> kafkaBundles = new ArrayList<>(); final String kafkaId = "nifi-kafka-0-10-nar"; kafkaBundles.add(generateBundle(defaultGroup, kafkaId, anotherVersion, standardSvcsBundles.get(1).getBundleDetails().getCoordinate())); Mockito.when(minifi.getBundles(anyString())).thenAnswer((Answer<List<Bundle>>) invocationOnMock -> { final String requestedClass = (String) invocationOnMock.getArguments()[0]; switch (requestedClass) { case KAFKA_PROCESSOR_CLASS: return kafkaBundles; case SSL_CONTROLLER_SERVICE_CLASS: return sslBundles; default: return minifiStdBundles; } }); final FlowEnricher flowEnricher = new FlowEnricher(minifi, mockFlowParser, mockProperties); flowEnricher.enrichFlowWithBundleInformation(); final ArgumentCaptor<String> bundleLookupCaptor = ArgumentCaptor.forClass(String.class); Mockito.verify(minifi, atLeastOnce()).getBundles(bundleLookupCaptor.capture()); final List<String> allValues = bundleLookupCaptor.getAllValues(); EXTENSION_CLASSES.stream().forEach(procClass -> Assert.assertTrue(allValues.contains(procClass))); final NodeList processorNodes = flowDocument.getElementsByTagName(FlowEnricher.PROCESSOR_TAG_NAME); Assert.assertEquals("Incorrect number of processors", 3, processorNodes.getLength()); final NodeList controllerServiceNodes = flowDocument.getElementsByTagName(FlowEnricher.CONTROLLER_SERVICE_TAG_NAME); Assert.assertEquals("Incorrect number of controller services", 1, controllerServiceNodes.getLength()); final NodeList reportingTaskNodes = flowDocument.getElementsByTagName(FlowEnricher.REPORTING_TASK_TAG_NAME); Assert.assertEquals("Incorrect number of reporting tasks", 0, reportingTaskNodes.getLength()); for (int i = 0; i < processorNodes.getLength(); i++) { final Element componentElement = (Element) processorNodes.item(i); Assert.assertEquals(1, componentElement.getElementsByTagName(FlowEnricher.EnrichingElementAdapter.BUNDLE_ELEMENT_NAME).getLength()); final FlowEnricher.EnrichingElementAdapter elementAdapter = new FlowEnricher.EnrichingElementAdapter(componentElement); if (elementAdapter.getComponentClass().equalsIgnoreCase(KAFKA_PROCESSOR_CLASS)) { verifyBundleProperties(elementAdapter, defaultGroup, kafkaId, anotherVersion); } else { verifyBundleProperties(elementAdapter, minifiGroup, minifiStandardNarId, minifiVersion); } } for (int i = 0; i < controllerServiceNodes.getLength(); i++) { final Element componentElement = (Element) controllerServiceNodes.item(i); Assert.assertEquals(1, componentElement.getElementsByTagName(FlowEnricher.EnrichingElementAdapter.BUNDLE_ELEMENT_NAME).getLength()); final FlowEnricher.EnrichingElementAdapter elementAdapter = new FlowEnricher.EnrichingElementAdapter(componentElement); if (elementAdapter.getComponentClass().equalsIgnoreCase(SSL_CONTROLLER_SERVICE_CLASS)) { verifyBundleProperties(elementAdapter, defaultGroup, sslContextSvcId, anotherVersion); } else { Assert.fail("A controller service that was not an SSL Controller service was found."); } } Mockito.verify(mockFlowParser, atLeastOnce()).writeFlow(any(), any()); }