azkaban.execapp.FlowRunnerTestUtil.java Source code

Java tutorial

Introduction

Here is the source code for azkaban.execapp.FlowRunnerTestUtil.java

Source

/*
 * Copyright 2017 LinkedIn Corp.
 *
 * Licensed 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 azkaban.execapp;

import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import azkaban.event.Event;
import azkaban.execapp.event.FlowWatcher;
import azkaban.execapp.jmx.JmxJobMBeanManager;
import azkaban.executor.ExecutableFlow;
import azkaban.executor.ExecutionOptions;
import azkaban.executor.ExecutionOptions.FailureAction;
import azkaban.executor.ExecutorLoader;
import azkaban.executor.InteractiveTestJob;
import azkaban.executor.MockExecutorLoader;
import azkaban.executor.Status;
import azkaban.flow.Flow;
import azkaban.jobtype.JobTypeManager;
import azkaban.jobtype.JobTypePluginSet;
import azkaban.project.FlowLoader;
import azkaban.project.FlowLoaderFactory;
import azkaban.project.Project;
import azkaban.project.ProjectLoader;
import azkaban.project.ProjectManagerException;
import azkaban.test.Utils;
import azkaban.test.executions.ExecutionsTestUtil;
import azkaban.utils.JSONUtils;
import azkaban.utils.Props;
import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import org.apache.commons.io.FileUtils;
import org.junit.rules.TemporaryFolder;

public class FlowRunnerTestUtil {

    private static int id = 101;
    private final Map<String, Flow> flowMap;
    private final Project project;
    private final File workingDir;
    private final JobTypeManager jobtypeManager;
    private final File projectDir;
    private final ProjectLoader projectLoader;
    private ExecutorLoader executorLoader;

    public FlowRunnerTestUtil(final String flowName, final TemporaryFolder temporaryFolder) throws Exception {

        this.projectDir = ExecutionsTestUtil.getFlowDir(flowName);
        this.workingDir = temporaryFolder.newFolder();
        this.project = new Project(1, "testProject");

        this.flowMap = FlowRunnerTestUtil.prepareProject(this.project, this.projectDir, this.workingDir);

        this.executorLoader = mock(ExecutorLoader.class);
        this.projectLoader = mock(ProjectLoader.class);
        when(this.executorLoader.updateExecutableReference(anyInt(), anyLong())).thenReturn(true);

        Utils.initServiceProvider();
        JmxJobMBeanManager.getInstance().initialize(new Props());

        InteractiveTestJob.clearTestJobs();

        this.jobtypeManager = new JobTypeManager(null, null, this.getClass().getClassLoader());
        final JobTypePluginSet pluginSet = this.jobtypeManager.getJobTypePluginSet();
        pluginSet.addPluginClass("test", InteractiveTestJob.class);
    }

    /**
     * Initialize the project with the flow definitions stored in the given source directory. Also
     * copy the source directory to the working directory.
     *
     * @param project project to initialize
     * @param sourceDir the source dir
     * @param workingDir the working dir
     * @return the flow name to flow map
     * @throws ProjectManagerException the project manager exception
     * @throws IOException the io exception
     */
    public static Map<String, Flow> prepareProject(final Project project, final File sourceDir,
            final File workingDir) throws ProjectManagerException, IOException {
        final FlowLoaderFactory loaderFactory = new FlowLoaderFactory(new Props(null));
        final FlowLoader loader = loaderFactory.createFlowLoader(sourceDir);

        loader.loadProjectFlow(project, sourceDir);
        if (!loader.getErrors().isEmpty()) {
            for (final String error : loader.getErrors()) {
                System.out.println(error);
            }
            throw new RuntimeException(String.format(
                    "Errors found in loading flows into a project ( %s ). From the directory: ( %s ).",
                    project.getName(), sourceDir));
        }

        final Map<String, Flow> flowMap = loader.getFlowMap();
        project.setFlows(flowMap);
        FileUtils.copyDirectory(sourceDir, workingDir);
        return flowMap;
    }

    public static void waitEventFired(final EventCollectorListener eventCollector, final String nestedId,
            final Status status) throws InterruptedException {
        for (int i = 0; i < 1000; i++) {
            for (final Event event : eventCollector.getEventList()) {
                if (event.getData().getStatus() == status && event.getData().getNestedId().equals(nestedId)) {
                    return;
                }
            }
            synchronized (EventCollectorListener.handleEvent) {
                EventCollectorListener.handleEvent.wait(10L);
            }
        }
        fail("Event wasn't fired with [" + nestedId + "], " + status);
    }

    public static ExecutableFlow prepareExecDir(final File workingDir, final File execDir, final String flowName,
            final int execId) throws IOException {
        FileUtils.copyDirectory(execDir, workingDir);
        final File jsonFlowFile = new File(workingDir, flowName + ".flow");
        final Object flowObj = JSONUtils.parseJSONFromFile(jsonFlowFile);
        final Project project = new Project(1, "test");
        final Flow flow = Flow.flowFromObject(flowObj);
        final ExecutableFlow execFlow = new ExecutableFlow(project, flow);
        execFlow.setExecutionId(execId);
        execFlow.setExecutionPath(workingDir.getPath());
        return execFlow;
    }

    public static void startThread(final FlowRunner runner) {
        new Thread(runner).start();
    }

    public FlowRunner createFromFlowFile(final String flowName) throws Exception {
        return createFromFlowFile(new EventCollectorListener(), flowName);
    }

    public FlowRunner createFromFlowFile(final EventCollectorListener eventCollector, final String flowName)
            throws Exception {
        return createFromFlowFile(flowName, eventCollector, new ExecutionOptions());
    }

    public FlowRunner createFromFlowFile(final String flowName, final EventCollectorListener eventCollector,
            final ExecutionOptions options) throws Exception {
        return createFromFlowFile(flowName, eventCollector, options, null, null);
    }

    public FlowRunner createFromFlowFile(final String flowName, final FlowWatcher watcher, final Integer pipeline)
            throws Exception {
        return createFromFlowFile(flowName, new EventCollectorListener(), new ExecutionOptions(), watcher,
                pipeline);
    }

    public FlowRunner createFromFlowFile(final String flowName, final EventCollectorListener eventCollector,
            final ExecutionOptions options, final FlowWatcher watcher, final Integer pipeline) throws Exception {
        final ExecutableFlow exFlow = FlowRunnerTestUtil.prepareExecDir(this.workingDir, this.projectDir, flowName,
                1);
        if (watcher != null) {
            options.setPipelineLevel(pipeline);
            options.setPipelineExecutionId(watcher.getExecId());
        }
        final FlowRunner runner = createFromExecutableFlow(eventCollector, exFlow, options, new HashMap<>(),
                new Props());
        runner.setFlowWatcher(watcher);
        return runner;
    }

    public FlowRunner createFromFlowMap(final String flowName, final String jobIdPrefix) throws Exception {
        return createFromFlowMap(flowName, jobIdPrefix, new ExecutionOptions());
    }

    public FlowRunner createFromFlowMap(final String flowName, final String jobIdPrefix,
            final ExecutionOptions options) throws Exception {
        return createFromFlowMap(flowName, jobIdPrefix, options, new Props());
    }

    public FlowRunner createFromFlowMap(final String flowName, final HashMap<String, String> flowParams)
            throws Exception {
        return createFromFlowMap(null, flowName, null, flowParams, new Props());
    }

    public FlowRunner createFromFlowMap(final String flowName, final ExecutionOptions options,
            final Map<String, String> flowParams, final Props azkabanProps) throws Exception {
        return createFromFlowMap(null, flowName, options, flowParams, azkabanProps);
    }

    public FlowRunner createFromFlowMap(final EventCollectorListener eventCollector, final String flowName,
            final ExecutionOptions options, final Map<String, String> flowParams, final Props azkabanProps)
            throws Exception {
        final Flow flow = this.flowMap.get(flowName);
        final ExecutableFlow exFlow = new ExecutableFlow(this.project, flow);
        return createFromExecutableFlow(eventCollector, exFlow, options, flowParams, azkabanProps);
    }

    public FlowRunner createFromFlowMap(final String flowName, final FailureAction action) throws Exception {
        final ExecutionOptions options = new ExecutionOptions();
        options.setFailureAction(action);
        return createFromFlowMap(flowName, options, new HashMap<>(), new Props());
    }

    private FlowRunner createFromFlowMap(final String flowName, final String jobIdPrefix,
            final ExecutionOptions options, final Props azkabanProps) throws Exception {
        final Map<String, String> flowParams = new HashMap<>();
        flowParams.put(InteractiveTestJob.JOB_ID_PREFIX, jobIdPrefix);
        return createFromFlowMap(new EventCollectorListener(), flowName, options, flowParams, azkabanProps);
    }

    private FlowRunner createFromExecutableFlow(final EventCollectorListener eventCollector,
            final ExecutableFlow exFlow, final ExecutionOptions options, final Map<String, String> flowParams,
            final Props azkabanProps) throws Exception {
        final int exId = id++;
        exFlow.setExecutionPath(this.workingDir.getPath());
        exFlow.setExecutionId(exId);
        if (options != null) {
            exFlow.setExecutionOptions(options);
        }
        exFlow.getExecutionOptions().addAllFlowParameters(flowParams);
        this.executorLoader.uploadExecutableFlow(exFlow);
        final FlowRunner runner = new FlowRunner(exFlow, this.executorLoader, this.projectLoader,
                this.jobtypeManager, azkabanProps, null);
        if (eventCollector != null) {
            runner.addListener(eventCollector);
        }
        return runner;
    }

    public ExecutorLoader getExecutorLoader() {
        return this.executorLoader;
    }

    public void setExecutorLoader(final MockExecutorLoader executorLoader) {
        this.executorLoader = executorLoader;
    }

    public ProjectLoader getProjectLoader() {
        return this.projectLoader;
    }

    public Project getProject() {
        return this.project;
    }
}