org.apache.nifi.reporting.ReportingInitializationContext Java Examples

private void initContexts() {
    configurationContext = Mockito.mock(ConfigurationContext.class);
    context = Mockito.mock(ReportingContext.class);
    Mockito.when(context.getProperty(DataDogReportingTask.ENVIRONMENT))
            .thenReturn(new MockPropertyValue(env, null));
    Mockito.when(context.getProperty(DataDogReportingTask.METRICS_PREFIX))
            .thenReturn(new MockPropertyValue(prefix, null));
    Mockito.when(context.getProperty(DataDogReportingTask.API_KEY))
            .thenReturn(new MockPropertyValue("agent", null));
    Mockito.when(context.getProperty(DataDogReportingTask.DATADOG_TRANSPORT))
            .thenReturn(new MockPropertyValue("DataDog Agent", null));
    EventAccess eventAccess = Mockito.mock(EventAccess.class);
    Mockito.when(eventAccess.getControllerStatus()).thenReturn(status);
    Mockito.when(context.getEventAccess()).thenReturn(eventAccess);

    logger = Mockito.mock(Logger.class);
    initContext = Mockito.mock(ReportingInitializationContext.class);
    Mockito.when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    //Mockito.when(initContext.getLogger()).thenReturn(logger);
    metricsMap = new ConcurrentHashMap<>();
    metricRegistry = Mockito.mock(MetricRegistry.class);
    virtualMachineMetrics = VirtualMachineMetrics.getInstance();
    metricsService = Mockito.mock(MetricsService.class);

}
 

private void initContexts() {
    configurationContext = Mockito.mock(ConfigurationContext.class);
    context = Mockito.mock(ReportingContext.class);
    Mockito.when(context.getProperty(DataDogReportingTask.ENVIRONMENT))
            .thenReturn(new MockPropertyValue(env, null));
    Mockito.when(context.getProperty(DataDogReportingTask.METRICS_PREFIX))
            .thenReturn(new MockPropertyValue(prefix, null));
    Mockito.when(context.getProperty(DataDogReportingTask.API_KEY))
            .thenReturn(new MockPropertyValue("agent", null));
    Mockito.when(context.getProperty(DataDogReportingTask.DATADOG_TRANSPORT))
            .thenReturn(new MockPropertyValue("DataDog Agent", null));
    EventAccess eventAccess = Mockito.mock(EventAccess.class);
    Mockito.when(eventAccess.getControllerStatus()).thenReturn(status);
    Mockito.when(context.getEventAccess()).thenReturn(eventAccess);

    logger = Mockito.mock(Logger.class);
    initContext = Mockito.mock(ReportingInitializationContext.class);
    Mockito.when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    //Mockito.when(initContext.getLogger()).thenReturn(logger);
    metricsMap = new ConcurrentHashMap<>();
    metricRegistry = Mockito.mock(MetricRegistry.class);
    virtualMachineMetrics = JmxJvmMetrics.getInstance();
    metricsService = Mockito.mock(MetricsService.class);

}
 

@Before
public void setup() throws InitializationException {
    System.setProperty(NiFiProperties.PROPERTIES_FILE_PATH, TestStandardProcessScheduler.class.getResource("/nifi.properties").getFile());
    this.nifiProperties = NiFiProperties.createBasicNiFiProperties(null, null);
    scheduler = new StandardProcessScheduler(Mockito.mock(ControllerServiceProvider.class), null, stateMgrProvider, variableRegistry, nifiProperties);
    scheduler.setSchedulingAgent(SchedulingStrategy.TIMER_DRIVEN, Mockito.mock(SchedulingAgent.class));

    reportingTask = new TestReportingTask();
    final ReportingInitializationContext config = new StandardReportingInitializationContext(UUID.randomUUID().toString(), "Test", SchedulingStrategy.TIMER_DRIVEN, "5 secs",
            Mockito.mock(ComponentLog.class), null, nifiProperties);
    reportingTask.initialize(config);

    final ValidationContextFactory validationContextFactory = new StandardValidationContextFactory(null, variableRegistry);
    final ComponentLog logger = Mockito.mock(ComponentLog.class);
    taskNode = new StandardReportingTaskNode(reportingTask, UUID.randomUUID().toString(), null, scheduler, validationContextFactory, variableRegistry, logger);

    controller = Mockito.mock(FlowController.class);
    rootGroup = new MockProcessGroup();
    Mockito.when(controller.getGroup(Mockito.anyString())).thenReturn(rootGroup);
}
 

@Override
public void initialize(ConfigurableComponent component) throws InitializationException {
    ReportingTask reportingTask = (ReportingTask) component;
    ReportingInitializationContext context = new MockReportingInitializationContext();
    try (NarCloseable narCloseable = NarCloseable.withComponentNarLoader(extensionManager, component.getClass(), context.getIdentifier())) {
        reportingTask.initialize(context);
    }
}
 

private LoggableComponent<ReportingTask> createLoggableReportingTask() throws ReportingTaskInstantiationException {
    try {
        final LoggableComponent<ReportingTask> taskComponent = createLoggableComponent(ReportingTask.class);

        final String taskName = taskComponent.getComponent().getClass().getSimpleName();
        final ReportingInitializationContext config = new StandardReportingInitializationContext(identifier, taskName,
                SchedulingStrategy.TIMER_DRIVEN, "1 min", taskComponent.getLogger(), serviceProvider, kerberosConfig, nodeTypeProvider);

        taskComponent.getComponent().initialize(config);

        return taskComponent;
    } catch (final Exception e) {
        throw new ReportingTaskInstantiationException(type, e);
    }
}
 

@Override
protected void init(final ReportingInitializationContext config) {
    metricsQueryService = new MetricsSqlQueryService(getLogger());
    final List<PropertyDescriptor> properties = new ArrayList<>();
    properties.add(QueryMetricsUtil.QUERY);
    properties.add(QueryMetricsUtil.RECORD_SINK);
    properties.add(QueryMetricsUtil.INCLUDE_ZERO_RECORD_RESULTS);
    this.properties = Collections.unmodifiableList(properties);
}
 

@Override
protected void init(final ReportingInitializationContext config) {
    metricsQueryService = new MetricsSqlQueryService(getLogger());
    final List<PropertyDescriptor> properties = new ArrayList<>();
    properties.add(QueryMetricsUtil.QUERY);
    properties.add(QueryMetricsUtil.RULES_ENGINE);
    properties.add(QueryMetricsUtil.ACTION_HANDLER);
    this.properties = Collections.unmodifiableList(properties);
}
 

/**
 * Register all wanted metrics to {@link #metricRegistry}.
 * <p>
 * {@inheritDoc}
 */
@Override
protected void init(ReportingInitializationContext config) {
    metricRegistry = new MetricRegistry();
    currentStatusReference = new AtomicReference<>();
    metricRegistry.registerAll(new MemoryUsageGaugeSet());
    metricRegistry.registerAll(new FlowMetricSet(currentStatusReference));
}
 

@Before
public void setUp() throws Exception {
    testSubject = new ReportLineageToAtlas();
    componentLogger = new MockComponentLog("reporting-task-id", testSubject);

    initializationContext = mock(ReportingInitializationContext.class);
    when(initializationContext.getLogger()).thenReturn(componentLogger);
}
 

@Override
public void initialize(ConfigurableComponent component) throws InitializationException {
    ReportingTask reportingTask = (ReportingTask) component;
    ReportingInitializationContext context = new MockReportingInitializationContext();
    try (NarCloseable narCloseable = NarCloseable.withComponentNarLoader(component.getClass(), context.getIdentifier())) {
        reportingTask.initialize(context);
    }
}
 

private ReportingTaskNode getOrCreateReportingTask(final FlowController controller, final ReportingTaskDTO dto, final boolean controllerInitialized, final boolean existingFlowEmpty)
        throws ReportingTaskInstantiationException {
    // create a new reporting task node when the controller is not initialized or the flow is empty
    if (!controllerInitialized || existingFlowEmpty) {
        final ReportingTaskNode reportingTask = controller.createReportingTask(dto.getType(), dto.getId(), false);
        reportingTask.setName(dto.getName());
        reportingTask.setComments(dto.getComments());
        reportingTask.setSchedulingPeriod(dto.getSchedulingPeriod());
        reportingTask.setSchedulingStrategy(SchedulingStrategy.valueOf(dto.getSchedulingStrategy()));

        reportingTask.setAnnotationData(dto.getAnnotationData());
        reportingTask.setProperties(dto.getProperties());

        final ComponentLog componentLog = new SimpleProcessLogger(dto.getId(), reportingTask.getReportingTask());
        final ReportingInitializationContext config = new StandardReportingInitializationContext(dto.getId(), dto.getName(),
                SchedulingStrategy.valueOf(dto.getSchedulingStrategy()), dto.getSchedulingPeriod(), componentLog, controller, nifiProperties);

        try {
            reportingTask.getReportingTask().initialize(config);
        } catch (final InitializationException ie) {
            throw new ReportingTaskInstantiationException("Failed to initialize reporting task of type " + dto.getType(), ie);
        }

        return reportingTask;
    } else {
        // otherwise return the existing reporting task node
        return controller.getReportingTaskNode(dto.getId());
    }
}
 

@Override
public void initialize(ConfigurableComponent component) throws InitializationException {
    ReportingTask reportingTask = (ReportingTask) component;
    ReportingInitializationContext context = new MockReportingInitializationContext();
    try (NarCloseable narCloseable = NarCloseable.withComponentNarLoader(component.getClass(), context.getIdentifier())) {
        reportingTask.initialize(context);
    }
}
 

private void test(TestConfiguration tc) throws Exception {
    final ReportLineageToAtlas reportingTask = new ReportLineageToAtlas();
    final MockComponentLog logger = new MockComponentLog("reporting-task-id", reportingTask);

    final ReportingInitializationContext initializationContext = mock(ReportingInitializationContext.class);
    when(initializationContext.getLogger()).thenReturn(logger);
    final ConfigurationContext configurationContext = new MockConfigurationContext(tc.properties, null);
    final ValidationContext validationContext = mock(ValidationContext.class);
    when(validationContext.getProperty(any())).then(invocation -> new MockPropertyValue(tc.properties.get(invocation.getArguments()[0])));
    final ReportingContext reportingContext = mock(ReportingContext.class);
    final MockStateManager stateManager = new MockStateManager(reportingTask);
    final EventAccess eventAccess = mock(EventAccess.class);
    when(reportingContext.getProperties()).thenReturn(tc.properties);
    when(reportingContext.getProperty(any())).then(invocation -> new MockPropertyValue(tc.properties.get(invocation.getArguments()[0])));
    when(reportingContext.getStateManager()).thenReturn(stateManager);
    when(reportingContext.getEventAccess()).thenReturn(eventAccess);
    when(eventAccess.getGroupStatus(eq("root"))).thenReturn(tc.rootPgStatus);

    final ProvenanceRepository provenanceRepository = mock(ProvenanceRepository.class);
    when(eventAccess.getControllerStatus()).thenReturn(tc.rootPgStatus);
    when(eventAccess.getProvenanceRepository()).thenReturn(provenanceRepository);
    when(eventAccess.getProvenanceEvents(eq(-1L), anyInt())).thenReturn(tc.provenanceRecords);
    when(provenanceRepository.getMaxEventId()).thenReturn((long) tc.provenanceRecords.size() - 1);
    when(provenanceRepository.getEvent(anyLong())).then(invocation -> tc.provenanceRecords.get(((Long) invocation.getArguments()[0]).intValue()));

    // To mock this async method invocations, keep the requested event ids in a stack.
    final ComputeLineageSubmission lineageComputationSubmission = mock(ComputeLineageSubmission.class);
    when(provenanceRepository.submitLineageComputation(anyLong(), any())).thenAnswer(invocation -> {
        requestedLineageComputationIds.push((Long) invocation.getArguments()[0]);
        return lineageComputationSubmission;
    });
    when(lineageComputationSubmission.getResult()).then(invocation -> tc.lineageResults.get(requestedLineageComputationIds.pop()));

    final ComputeLineageSubmission expandParentsSubmission = mock(ComputeLineageSubmission.class);
    when(provenanceRepository.submitExpandParents(anyLong(), any())).thenAnswer(invocation -> {
        requestedExpandParentsIds.push(((Long) invocation.getArguments()[0]));
        return expandParentsSubmission;
    });
    when(expandParentsSubmission.getResult()).then(invocation -> tc.parentLineageResults.get(requestedExpandParentsIds.pop()));

    tc.properties.put(ATLAS_NIFI_URL, "http://localhost:8080/nifi");
    tc.properties.put(ATLAS_URLS, TARGET_ATLAS_URL);
    tc.properties.put(ATLAS_USER, "admin");
    tc.properties.put(ATLAS_PASSWORD, "admin");
    tc.properties.put(new PropertyDescriptor.Builder().name("hostnamePattern.example").dynamic(true).build(), ".*");


    reportingTask.initialize(initializationContext);
    reportingTask.validate(validationContext);
    reportingTask.setup(configurationContext);
    reportingTask.onTrigger(reportingContext);
    reportingTask.onUnscheduled();
    reportingTask.onStopped();
}
 

@Test
public void testOnTrigger() throws InitializationException {
    final String metricsUrl = "http://localhost:9091";
    final String applicationId = "nifi";
    final String hostName = "localhost";
    final String jobName = "nifi_reporting_job";
    final boolean jvmMetrics = true;
    final boolean authentication = false;

    // create the jersey client mocks for handling the post
    final Client client = Mockito.mock(Client.class);
    final WebTarget target = Mockito.mock(WebTarget.class);
    final Invocation.Builder builder = Mockito.mock(Invocation.Builder.class);

    final Response response = Mockito.mock(Response.class);
    Mockito.when(response.getStatus()).thenReturn(200);

    Mockito.when(client.target(metricsUrl)).thenReturn(target);
    Mockito.when(target.request()).thenReturn(builder);
    Mockito.when(builder.post(Matchers.any(Entity.class))).thenReturn(response);

    // mock the ReportingInitializationContext for initialize(...)
    final ComponentLog logger = Mockito.mock(ComponentLog.class);
    final ReportingInitializationContext initContext = Mockito.mock(ReportingInitializationContext.class);
    Mockito.when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    Mockito.when(initContext.getLogger()).thenReturn(logger);


    // mock the ReportingContext for onTrigger(...)
    final ReportingContext context = Mockito.mock(ReportingContext.class);
    Mockito.when(context.getProperty(PrometheusReportingTask.METRICS_COLLECTOR_URL))
            .thenReturn(new MockPropertyValue(metricsUrl));
    Mockito.when(context.getProperty(PrometheusReportingTask.APPLICATION_ID))
            .thenReturn(new MockPropertyValue(applicationId));
    Mockito.when(context.getProperty(PrometheusReportingTask.INSTANCE_ID))
            .thenReturn(new MockPropertyValue(hostName));
    Mockito.when(context.getProperty(PrometheusReportingTask.PROCESS_GROUP_IDS))
            .thenReturn(new MockPropertyValue("1234"));
    Mockito.when(context.getProperty(PrometheusReportingTask.JOB_NAME))
            .thenReturn(new MockPropertyValue(jobName));
    Mockito.when(context.getProperty(PrometheusReportingTask.SEND_JVM_METRICS))
            .thenReturn(new MockPropertyValue(Boolean.toString(jvmMetrics)));
    Mockito.when(context.getProperty(PrometheusReportingTask.USE_AUTHENTICATION))
            .thenReturn(new MockPropertyValue(Boolean.toString(authentication)));

    final EventAccess eventAccess = Mockito.mock(EventAccess.class);
    Mockito.when(context.getEventAccess()).thenReturn(eventAccess);
    Mockito.when(eventAccess.getControllerStatus()).thenReturn(status);
    Mockito.when(eventAccess.getGroupStatus("1234")).thenReturn(status);

    // create a testable instance of the reporting task
    final PrometheusReportingTask task = new PrometheusReportingTask();
    task.initialize(initContext);
    task.onTrigger(context);
}
 

@Test
public void testOnTrigger() throws InitializationException, IOException {
    final String metricsUrl = "http://myambari:6188/ws/v1/timeline/metrics";
    final String applicationId = "NIFI";
    final String hostName = "localhost";

    // create the jersey client mocks for handling the post
    final Client client = Mockito.mock(Client.class);
    final WebTarget target = Mockito.mock(WebTarget.class);
    final Invocation.Builder builder = Mockito.mock(Invocation.Builder.class);

    final Response response = Mockito.mock(Response.class);
    Mockito.when(response.getStatus()).thenReturn(200);

    Mockito.when(client.target(metricsUrl)).thenReturn(target);
    Mockito.when(target.request()).thenReturn(builder);
    Mockito.when(builder.post(ArgumentMatchers.any(Entity.class))).thenReturn(response);

    // mock the ReportingInitializationContext for initialize(...)
    final ComponentLog logger = Mockito.mock(ComponentLog.class);
    final ReportingInitializationContext initContext = Mockito.mock(ReportingInitializationContext.class);
    Mockito.when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    Mockito.when(initContext.getLogger()).thenReturn(logger);

    // mock the ConfigurationContext for setup(...)
    final ConfigurationContext configurationContext = Mockito.mock(ConfigurationContext.class);

    // mock the ReportingContext for onTrigger(...)
    final ReportingContext context = Mockito.mock(ReportingContext.class);
    Mockito.when(context.getProperty(AmbariReportingTask.METRICS_COLLECTOR_URL))
            .thenReturn(new MockPropertyValue(metricsUrl));
    Mockito.when(context.getProperty(AmbariReportingTask.APPLICATION_ID))
            .thenReturn(new MockPropertyValue(applicationId));
    Mockito.when(context.getProperty(AmbariReportingTask.HOSTNAME))
            .thenReturn(new MockPropertyValue(hostName));
    Mockito.when(context.getProperty(AmbariReportingTask.PROCESS_GROUP_ID))
            .thenReturn(new MockPropertyValue("1234"));


    final EventAccess eventAccess = Mockito.mock(EventAccess.class);
    Mockito.when(context.getEventAccess()).thenReturn(eventAccess);
    Mockito.when(eventAccess.getControllerStatus()).thenReturn(status);

    // create a testable instance of the reporting task
    final AmbariReportingTask task = new TestableAmbariReportingTask(client);
    task.initialize(initContext);
    task.setup(configurationContext);
    task.onTrigger(context);
}
 

@Override
public void init(ReportingInitializationContext config)
        throws InitializationException {
   config.getLogger().info("Initializing...");
}
 

@Override
public void init(ReportingInitializationContext config)
        throws InitializationException {
   config.getLogger().info("Initializing...");
}
 

@Test
public void testOnTrigger() throws InitializationException, IOException {
    final String metricsUrl = "http://myambari:6188/ws/v1/timeline/metrics";
    final String applicationId = "NIFI";
    final String hostName = "localhost";

    // create the jersey client mocks for handling the post
    final Client client = Mockito.mock(Client.class);
    final WebTarget target = Mockito.mock(WebTarget.class);
    final Invocation.Builder builder = Mockito.mock(Invocation.Builder.class);

    final Response response = Mockito.mock(Response.class);
    Mockito.when(response.getStatus()).thenReturn(200);

    Mockito.when(client.target(metricsUrl)).thenReturn(target);
    Mockito.when(target.request()).thenReturn(builder);
    Mockito.when(builder.post(Matchers.any(Entity.class))).thenReturn(response);

    // mock the ReportingInitializationContext for initialize(...)
    final ComponentLog logger = Mockito.mock(ComponentLog.class);
    final ReportingInitializationContext initContext = Mockito.mock(ReportingInitializationContext.class);
    Mockito.when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    Mockito.when(initContext.getLogger()).thenReturn(logger);

    // mock the ConfigurationContext for setup(...)
    final ConfigurationContext configurationContext = Mockito.mock(ConfigurationContext.class);

    // mock the ReportingContext for onTrigger(...)
    final ReportingContext context = Mockito.mock(ReportingContext.class);
    Mockito.when(context.getProperty(AmbariReportingTask.METRICS_COLLECTOR_URL))
            .thenReturn(new MockPropertyValue(metricsUrl));
    Mockito.when(context.getProperty(AmbariReportingTask.APPLICATION_ID))
            .thenReturn(new MockPropertyValue(applicationId));
    Mockito.when(context.getProperty(AmbariReportingTask.HOSTNAME))
            .thenReturn(new MockPropertyValue(hostName));
    Mockito.when(context.getProperty(AmbariReportingTask.PROCESS_GROUP_ID))
            .thenReturn(new MockPropertyValue("1234"));


    final EventAccess eventAccess = Mockito.mock(EventAccess.class);
    Mockito.when(context.getEventAccess()).thenReturn(eventAccess);
    Mockito.when(eventAccess.getControllerStatus()).thenReturn(status);

    // create a testable instance of the reporting task
    final AmbariReportingTask task = new TestableAmbariReportingTask(client);
    task.initialize(initContext);
    task.setup(configurationContext);
    task.onTrigger(context);
}