SimpleFlowMapper implements FlowMapper { @Override public Optional<FlowUri> getFlowMapping(final String agentClass) throws FlowMapperException { if (StringUtils.isBlank(agentClass)) { throw new IllegalArgumentException("Agent class cannot be null or blank"); } try { final String registryUrl = clientFactory.getNiFiRegistryUrl(); final Optional<String> bucketIdOptional = clientFactory.resolveNiFiRegistryBucketId(); if (bucketIdOptional.isPresent()) { final String bucketId = bucketIdOptional.get(); final NiFiRegistryClient client = clientFactory.getClient(); final FlowClient flowClient = client.getFlowClient(); final List<VersionedFlow> flows = flowClient.getByBucket(bucketId); final Optional<String> flowId = flows.stream() .filter(f -> f.getName().equals(agentClass)) .map(f -> f.getIdentifier()) .findFirst(); if (flowId.isPresent()) { final FlowUri flowUri = new FlowUri(registryUrl, bucketId, flowId.get()); return Optional.of(flowUri); } } return Optional.empty(); } catch (final Exception e) { throw new FlowMapperException("Unable to get flow mapping for " + agentClass + " due to " + e.getMessage(), e); } } @Autowired  SimpleFlowMapper(final C2NiFiRegistryClient clientFactory); @Override Optional<FlowUri> getFlowMapping(final String agentClass);  }
@Test public void testGetFlowMappingWhenAgentClassExists() throws FlowMapperException, IOException, NiFiRegistryException { final String agentClass = "Class A"; final VersionedFlow versionedFlow = new VersionedFlow(); versionedFlow.setName(agentClass); versionedFlow.setIdentifier(UUID.randomUUID().toString()); versionedFlow.setBucketIdentifier(BUCKET_ID); when(flowClient.getByBucket(BUCKET_ID)).thenReturn(Collections.singletonList(versionedFlow)); final Optional<FlowUri> flowUri = flowMapper.getFlowMapping(agentClass); assertNotNull(flowUri); assertTrue(flowUri.isPresent()); assertEquals(REGISTRY_URL, flowUri.get().getRegistryUrl()); assertEquals(BUCKET_ID, flowUri.get().getBucketId()); assertEquals(versionedFlow.getIdentifier(), flowUri.get().getFlowId()); }
@Test public void testGetFlowMappingWhenAgentClassDoesNotExist() throws FlowMapperException, IOException, NiFiRegistryException { final String agentClass = "Class A"; final VersionedFlow versionedFlow = new VersionedFlow(); versionedFlow.setName("SOME-OTHER-FLOW"); versionedFlow.setIdentifier(UUID.randomUUID().toString()); versionedFlow.setBucketIdentifier(BUCKET_ID); when(flowClient.getByBucket(BUCKET_ID)).thenReturn(Collections.singletonList(versionedFlow)); final Optional<FlowUri> flowUri = flowMapper.getFlowMapping(agentClass); assertNotNull(flowUri); assertFalse(flowUri.isPresent()); }
@Test(expected = FlowMapperException.class) public void testGetFlowMappingWhenExceptionHappens() throws FlowMapperException, IOException, NiFiRegistryException { final String agentClass = "Class A"; when(flowClient.getByBucket(BUCKET_ID)).thenThrow(new NiFiRegistryException("Bucket not found")); flowMapper.getFlowMapping(agentClass); }