Browse Source
* Add the ProgramType parameter to distinguish task types * EmrAddStepsTask supports Add-Steps * UI supports Add-Steps * EmrTask modify the name of the class to EmrJobFlowTask * add ERM Task abstract base class AbstractEmrTask * add testcase for EmrAddStepsTask * Modifying help Documents3.1.0-release
ZhaoGuodong
2 years ago
committed by
GitHub
21 changed files with 769 additions and 125 deletions
After Width: | Height: | Size: 129 KiB |
After Width: | Height: | Size: 130 KiB |
After Width: | Height: | Size: 136 KiB |
@ -0,0 +1,113 @@ |
|||||||
|
/* |
||||||
|
* 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.dolphinscheduler.plugin.task.emr; |
||||||
|
|
||||||
|
import static com.fasterxml.jackson.databind.DeserializationFeature.ACCEPT_EMPTY_ARRAY_AS_NULL_OBJECT; |
||||||
|
import static com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES; |
||||||
|
import static com.fasterxml.jackson.databind.DeserializationFeature.READ_UNKNOWN_ENUM_VALUES_AS_NULL; |
||||||
|
import static com.fasterxml.jackson.databind.MapperFeature.REQUIRE_SETTERS_FOR_GETTERS; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.AbstractTaskExecutor; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskConstants; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.parameters.AbstractParameters; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.PropertyUtils; |
||||||
|
|
||||||
|
import java.util.TimeZone; |
||||||
|
|
||||||
|
import com.amazonaws.auth.AWSCredentialsProvider; |
||||||
|
import com.amazonaws.auth.AWSStaticCredentialsProvider; |
||||||
|
import com.amazonaws.auth.BasicAWSCredentials; |
||||||
|
import com.amazonaws.services.elasticmapreduce.AmazonElasticMapReduce; |
||||||
|
import com.amazonaws.services.elasticmapreduce.AmazonElasticMapReduceClientBuilder; |
||||||
|
import com.fasterxml.jackson.databind.ObjectMapper; |
||||||
|
import com.fasterxml.jackson.databind.PropertyNamingStrategy; |
||||||
|
|
||||||
|
/** |
||||||
|
* ERM Task abstract base class
|
||||||
|
* |
||||||
|
* @since v3.1.0 |
||||||
|
*/ |
||||||
|
public abstract class AbstractEmrTask extends AbstractTaskExecutor { |
||||||
|
|
||||||
|
final TaskExecutionContext taskExecutionContext; |
||||||
|
EmrParameters emrParameters; |
||||||
|
AmazonElasticMapReduce emrClient; |
||||||
|
String clusterId; |
||||||
|
|
||||||
|
|
||||||
|
/** |
||||||
|
* config ObjectMapper features and propertyNamingStrategy |
||||||
|
* use UpperCamelCaseStrategy support capital letters parse |
||||||
|
* |
||||||
|
* @see PropertyNamingStrategy.UpperCamelCaseStrategy |
||||||
|
*/ |
||||||
|
static final ObjectMapper objectMapper = new ObjectMapper() |
||||||
|
.configure(FAIL_ON_UNKNOWN_PROPERTIES, false) |
||||||
|
.configure(ACCEPT_EMPTY_ARRAY_AS_NULL_OBJECT, true) |
||||||
|
.configure(READ_UNKNOWN_ENUM_VALUES_AS_NULL, true) |
||||||
|
.configure(REQUIRE_SETTERS_FOR_GETTERS, true) |
||||||
|
.setTimeZone(TimeZone.getDefault()) |
||||||
|
.setPropertyNamingStrategy(new PropertyNamingStrategy.UpperCamelCaseStrategy()); |
||||||
|
|
||||||
|
/** |
||||||
|
* constructor |
||||||
|
* |
||||||
|
* @param taskExecutionContext taskExecutionContext |
||||||
|
*/ |
||||||
|
protected AbstractEmrTask(TaskExecutionContext taskExecutionContext) { |
||||||
|
super(taskExecutionContext); |
||||||
|
this.taskExecutionContext = taskExecutionContext; |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public void init() { |
||||||
|
final String taskParams = taskExecutionContext.getTaskParams(); |
||||||
|
logger.info("emr task params:{}", taskParams); |
||||||
|
emrParameters = JSONUtils.parseObject(taskParams, EmrParameters.class); |
||||||
|
if (emrParameters == null || !emrParameters.checkParameters()) { |
||||||
|
throw new EmrTaskException("emr task params is not valid"); |
||||||
|
} |
||||||
|
emrClient = createEmrClient(); |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public AbstractParameters getParameters() { |
||||||
|
return emrParameters; |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* create emr client from BasicAWSCredentials |
||||||
|
* |
||||||
|
* @return AmazonElasticMapReduce |
||||||
|
*/ |
||||||
|
private AmazonElasticMapReduce createEmrClient() { |
||||||
|
|
||||||
|
final String awsAccessKeyId = PropertyUtils.getString(TaskConstants.AWS_ACCESS_KEY_ID); |
||||||
|
final String awsSecretAccessKey = PropertyUtils.getString(TaskConstants.AWS_SECRET_ACCESS_KEY); |
||||||
|
final String awsRegion = PropertyUtils.getString(TaskConstants.AWS_REGION); |
||||||
|
final BasicAWSCredentials basicAWSCredentials = new BasicAWSCredentials(awsAccessKeyId, awsSecretAccessKey); |
||||||
|
final AWSCredentialsProvider awsCredentialsProvider = new AWSStaticCredentialsProvider(basicAWSCredentials); |
||||||
|
// create an EMR client
|
||||||
|
return AmazonElasticMapReduceClientBuilder.standard() |
||||||
|
.withCredentials(awsCredentialsProvider) |
||||||
|
.withRegion(awsRegion) |
||||||
|
.build(); |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,177 @@ |
|||||||
|
/* |
||||||
|
* 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.dolphinscheduler.plugin.task.emr; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskConstants; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext; |
||||||
|
|
||||||
|
import java.util.HashSet; |
||||||
|
import java.util.concurrent.TimeUnit; |
||||||
|
|
||||||
|
import com.amazonaws.SdkBaseException; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.AddJobFlowStepsRequest; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.AddJobFlowStepsResult; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.CancelStepsInfo; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.CancelStepsRequest; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.CancelStepsRequestStatus; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.CancelStepsResult; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.DescribeStepRequest; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.DescribeStepResult; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.StepState; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.StepStatus; |
||||||
|
import com.fasterxml.jackson.core.JsonProcessingException; |
||||||
|
import com.google.common.collect.Sets; |
||||||
|
|
||||||
|
/** |
||||||
|
* AddJobFlowSteps task executor |
||||||
|
* |
||||||
|
* @since v3.1.0 |
||||||
|
*/ |
||||||
|
public class EmrAddStepsTask extends AbstractEmrTask { |
||||||
|
|
||||||
|
private String stepId; |
||||||
|
|
||||||
|
private final HashSet<String> waitingStateSet = Sets.newHashSet( |
||||||
|
StepState.PENDING.toString(), |
||||||
|
StepState.CANCEL_PENDING.toString(), |
||||||
|
StepState.RUNNING.toString() |
||||||
|
); |
||||||
|
|
||||||
|
/** |
||||||
|
* constructor |
||||||
|
* |
||||||
|
* @param taskExecutionContext taskExecutionContext |
||||||
|
*/ |
||||||
|
protected EmrAddStepsTask(TaskExecutionContext taskExecutionContext) { |
||||||
|
super(taskExecutionContext); |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public void handle() throws InterruptedException { |
||||||
|
StepStatus stepStatus = null; |
||||||
|
try { |
||||||
|
AddJobFlowStepsRequest addJobFlowStepsRequest = createAddJobFlowStepsRequest(); |
||||||
|
|
||||||
|
// submit addJobFlowStepsRequest to aws
|
||||||
|
AddJobFlowStepsResult result = emrClient.addJobFlowSteps(addJobFlowStepsRequest); |
||||||
|
|
||||||
|
clusterId = addJobFlowStepsRequest.getJobFlowId(); |
||||||
|
stepId = result.getStepIds().get(0); |
||||||
|
// use clusterId-stepId as appIds
|
||||||
|
setAppIds(clusterId + TaskConstants.SUBTRACT_STRING + stepId); |
||||||
|
|
||||||
|
stepStatus = getStepStatus(); |
||||||
|
|
||||||
|
while (waitingStateSet.contains(stepStatus.getState())) { |
||||||
|
TimeUnit.SECONDS.sleep(10); |
||||||
|
stepStatus = getStepStatus(); |
||||||
|
} |
||||||
|
|
||||||
|
} catch (EmrTaskException | SdkBaseException e) { |
||||||
|
logger.error("emr task submit failed with error", e); |
||||||
|
} finally { |
||||||
|
final int exitStatusCode = calculateExitStatusCode(stepStatus); |
||||||
|
setExitStatusCode(exitStatusCode); |
||||||
|
logger.info("emr task finished with step status : {}", stepStatus); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* parse json string to AddJobFlowStepsRequest |
||||||
|
* |
||||||
|
* @return AddJobFlowStepsRequest |
||||||
|
*/ |
||||||
|
private AddJobFlowStepsRequest createAddJobFlowStepsRequest() { |
||||||
|
|
||||||
|
final AddJobFlowStepsRequest addJobFlowStepsRequest; |
||||||
|
try { |
||||||
|
addJobFlowStepsRequest = objectMapper.readValue(emrParameters.getStepsDefineJson(), AddJobFlowStepsRequest.class); |
||||||
|
} catch (JsonProcessingException e) { |
||||||
|
throw new EmrTaskException("can not parse AddJobFlowStepsRequest from json", e); |
||||||
|
} |
||||||
|
|
||||||
|
// When a single task definition is associated with multiple steps, the state tracking will have high complexity.
|
||||||
|
// Therefore, A task definition only supports the association of a single step, which can better ensure the reliability of the task state.
|
||||||
|
if (addJobFlowStepsRequest.getSteps().size() > 1) { |
||||||
|
throw new EmrTaskException("ds emr addJobFlowStepsTask only support one step"); |
||||||
|
} |
||||||
|
|
||||||
|
return addJobFlowStepsRequest; |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* calculate task exitStatusCode |
||||||
|
* |
||||||
|
* @param stepStatus aws emr execution status details of the cluster step. |
||||||
|
* @return exitStatusCode |
||||||
|
*/ |
||||||
|
private int calculateExitStatusCode(StepStatus stepStatus) { |
||||||
|
if (stepStatus == null) { |
||||||
|
return TaskConstants.EXIT_CODE_FAILURE; |
||||||
|
} else { |
||||||
|
String state = stepStatus.getState(); |
||||||
|
StepState stepState = StepState.valueOf(state); |
||||||
|
switch (stepState) { |
||||||
|
case COMPLETED: |
||||||
|
return TaskConstants.EXIT_CODE_SUCCESS; |
||||||
|
case CANCELLED: |
||||||
|
return TaskConstants.EXIT_CODE_KILL; |
||||||
|
default: |
||||||
|
return TaskConstants.EXIT_CODE_FAILURE; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
} |
||||||
|
|
||||||
|
private StepStatus getStepStatus() { |
||||||
|
DescribeStepRequest describeStepRequest = new DescribeStepRequest().withClusterId(clusterId).withStepId(stepId); |
||||||
|
DescribeStepResult result = emrClient.describeStep(describeStepRequest); |
||||||
|
if (result == null) { |
||||||
|
throw new EmrTaskException("fetch step status failed"); |
||||||
|
} |
||||||
|
StepStatus stepStatus = result.getStep().getStatus(); |
||||||
|
logger.info("emr step [clusterId:{}, stepId:{}] running with status:{}", clusterId, stepId, stepStatus); |
||||||
|
return stepStatus; |
||||||
|
|
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public void cancelApplication(boolean status) throws Exception { |
||||||
|
super.cancelApplication(status); |
||||||
|
logger.info("trying cancel emr step, taskId:{}, clusterId:{}, stepId:{}", this.taskExecutionContext.getTaskInstanceId(), clusterId, stepId); |
||||||
|
CancelStepsRequest cancelStepsRequest = new CancelStepsRequest().withClusterId(clusterId).withStepIds(stepId); |
||||||
|
CancelStepsResult cancelStepsResult = emrClient.cancelSteps(cancelStepsRequest); |
||||||
|
|
||||||
|
if (cancelStepsResult == null) { |
||||||
|
throw new EmrTaskException("cancel emr step failed"); |
||||||
|
} |
||||||
|
|
||||||
|
CancelStepsInfo cancelEmrStepInfo = cancelStepsResult.getCancelStepsInfoList() |
||||||
|
.stream() |
||||||
|
.filter(cancelStepsInfo -> cancelStepsInfo.getStepId().equals(stepId)) |
||||||
|
.findFirst() |
||||||
|
.orElseThrow(() -> new EmrTaskException("cancel emr step failed")); |
||||||
|
|
||||||
|
if (CancelStepsRequestStatus.FAILED.toString().equals(cancelEmrStepInfo.getStatus())) { |
||||||
|
throw new EmrTaskException("cancel emr step failed, message:" + cancelEmrStepInfo.getReason()); |
||||||
|
} |
||||||
|
|
||||||
|
logger.info("the result of cancel emr step is:{}", cancelStepsResult); |
||||||
|
} |
||||||
|
|
||||||
|
} |
@ -0,0 +1,34 @@ |
|||||||
|
/* |
||||||
|
* 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.dolphinscheduler.plugin.task.emr; |
||||||
|
|
||||||
|
/** |
||||||
|
* emr program type |
||||||
|
* |
||||||
|
* @since v3.1.0 |
||||||
|
*/ |
||||||
|
public enum ProgramType { |
||||||
|
/** |
||||||
|
* RunJobFlow |
||||||
|
*/ |
||||||
|
RUN_JOB_FLOW, |
||||||
|
/** |
||||||
|
* AddJobFlowSteps |
||||||
|
*/ |
||||||
|
ADD_JOB_FLOW_STEPS |
||||||
|
} |
@ -0,0 +1,198 @@ |
|||||||
|
/* |
||||||
|
* 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.dolphinscheduler.plugin.task.emr; |
||||||
|
|
||||||
|
import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_FAILURE; |
||||||
|
import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_KILL; |
||||||
|
import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_SUCCESS; |
||||||
|
|
||||||
|
import static org.mockito.ArgumentMatchers.any; |
||||||
|
import static org.powermock.api.mockito.PowerMockito.doReturn; |
||||||
|
import static org.powermock.api.mockito.PowerMockito.mock; |
||||||
|
import static org.powermock.api.mockito.PowerMockito.mockStatic; |
||||||
|
import static org.powermock.api.mockito.PowerMockito.spy; |
||||||
|
import static org.powermock.api.mockito.PowerMockito.when; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils; |
||||||
|
|
||||||
|
import org.apache.commons.io.IOUtils; |
||||||
|
|
||||||
|
import java.io.InputStream; |
||||||
|
import java.nio.charset.StandardCharsets; |
||||||
|
import java.util.Collections; |
||||||
|
|
||||||
|
import org.junit.Assert; |
||||||
|
import org.junit.Before; |
||||||
|
import org.junit.Test; |
||||||
|
import org.junit.runner.RunWith; |
||||||
|
import org.powermock.api.mockito.PowerMockito; |
||||||
|
import org.powermock.core.classloader.annotations.PowerMockIgnore; |
||||||
|
import org.powermock.core.classloader.annotations.PrepareForTest; |
||||||
|
import org.powermock.modules.junit4.PowerMockRunner; |
||||||
|
|
||||||
|
import com.amazonaws.services.elasticmapreduce.AmazonElasticMapReduce; |
||||||
|
import com.amazonaws.services.elasticmapreduce.AmazonElasticMapReduceClientBuilder; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.AddJobFlowStepsResult; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.AmazonElasticMapReduceException; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.DescribeStepResult; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.Step; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.StepState; |
||||||
|
import com.amazonaws.services.elasticmapreduce.model.StepStatus; |
||||||
|
|
||||||
|
/** |
||||||
|
* EmrAddStepsTask Test |
||||||
|
* |
||||||
|
* @since v3.1.0 |
||||||
|
*/ |
||||||
|
@RunWith(PowerMockRunner.class) |
||||||
|
@PrepareForTest({ |
||||||
|
AmazonElasticMapReduceClientBuilder.class, |
||||||
|
EmrAddStepsTask.class, |
||||||
|
AmazonElasticMapReduce.class, |
||||||
|
JSONUtils.class |
||||||
|
}) |
||||||
|
@PowerMockIgnore({"javax.*"}) |
||||||
|
public class EmrAddStepsTaskTest { |
||||||
|
|
||||||
|
private final StepStatus pendingState = |
||||||
|
new StepStatus().withState(StepState.PENDING); |
||||||
|
|
||||||
|
private final StepStatus runningState = |
||||||
|
new StepStatus().withState(StepState.RUNNING); |
||||||
|
|
||||||
|
private final StepStatus completedState = |
||||||
|
new StepStatus().withState(StepState.COMPLETED); |
||||||
|
|
||||||
|
private final StepStatus cancelledState = |
||||||
|
new StepStatus().withState(StepState.CANCELLED); |
||||||
|
|
||||||
|
private final StepStatus failedState = |
||||||
|
new StepStatus().withState(StepState.FAILED); |
||||||
|
|
||||||
|
private EmrAddStepsTask emrAddStepsTask; |
||||||
|
private AmazonElasticMapReduce emrClient; |
||||||
|
private Step step; |
||||||
|
|
||||||
|
@Before |
||||||
|
public void before() throws Exception { |
||||||
|
// mock EmrParameters and EmrAddStepsTask
|
||||||
|
EmrParameters emrParameters = buildEmrTaskParameters(); |
||||||
|
String emrParametersString = JSONUtils.toJsonString(emrParameters); |
||||||
|
TaskExecutionContext taskExecutionContext = PowerMockito.mock(TaskExecutionContext.class); |
||||||
|
when(taskExecutionContext.getTaskParams()).thenReturn(emrParametersString); |
||||||
|
emrAddStepsTask = spy(new EmrAddStepsTask(taskExecutionContext)); |
||||||
|
|
||||||
|
// mock emrClient and behavior
|
||||||
|
emrClient = mock(AmazonElasticMapReduce.class); |
||||||
|
|
||||||
|
AddJobFlowStepsResult addJobFlowStepsResult = mock(AddJobFlowStepsResult.class); |
||||||
|
when(emrClient.addJobFlowSteps(any())).thenReturn(addJobFlowStepsResult); |
||||||
|
when(addJobFlowStepsResult.getStepIds()).thenReturn(Collections.singletonList("step-xx")); |
||||||
|
|
||||||
|
doReturn(emrClient).when(emrAddStepsTask, "createEmrClient"); |
||||||
|
DescribeStepResult describeStepResult = mock(DescribeStepResult.class); |
||||||
|
when(emrClient.describeStep(any())).thenReturn(describeStepResult); |
||||||
|
|
||||||
|
// mock step
|
||||||
|
step = mock(Step.class); |
||||||
|
when(describeStepResult.getStep()).thenReturn(step); |
||||||
|
|
||||||
|
emrAddStepsTask.init(); |
||||||
|
} |
||||||
|
|
||||||
|
@Test |
||||||
|
public void testCanNotParseJson() throws Exception { |
||||||
|
mockStatic(JSONUtils.class); |
||||||
|
when(emrAddStepsTask, "createAddJobFlowStepsRequest").thenThrow(new EmrTaskException("can not parse AddJobFlowStepsRequest from json", new Exception("error"))); |
||||||
|
emrAddStepsTask.handle(); |
||||||
|
Assert.assertEquals(EXIT_CODE_FAILURE, emrAddStepsTask.getExitStatusCode()); |
||||||
|
} |
||||||
|
|
||||||
|
@Test |
||||||
|
public void testDefineJsonStepNotOne() throws Exception { |
||||||
|
// mock EmrParameters and EmrAddStepsTask
|
||||||
|
EmrParameters emrParameters = buildErrorEmrTaskParameters(); |
||||||
|
String emrParametersString = JSONUtils.toJsonString(emrParameters); |
||||||
|
TaskExecutionContext taskExecutionContext = PowerMockito.mock(TaskExecutionContext.class); |
||||||
|
when(taskExecutionContext.getTaskParams()).thenReturn(emrParametersString); |
||||||
|
emrAddStepsTask = spy(new EmrAddStepsTask(taskExecutionContext)); |
||||||
|
doReturn(emrClient).when(emrAddStepsTask, "createEmrClient"); |
||||||
|
emrAddStepsTask.init(); |
||||||
|
emrAddStepsTask.handle(); |
||||||
|
|
||||||
|
Assert.assertEquals(EXIT_CODE_FAILURE, emrAddStepsTask.getExitStatusCode()); |
||||||
|
} |
||||||
|
|
||||||
|
@Test |
||||||
|
public void testHandle() throws Exception { |
||||||
|
when(step.getStatus()).thenReturn(pendingState, runningState, completedState); |
||||||
|
|
||||||
|
emrAddStepsTask.handle(); |
||||||
|
Assert.assertEquals(EXIT_CODE_SUCCESS, emrAddStepsTask.getExitStatusCode()); |
||||||
|
} |
||||||
|
|
||||||
|
@Test |
||||||
|
public void testHandleUserRequestTerminate() throws Exception { |
||||||
|
when(step.getStatus()).thenReturn(pendingState, runningState, cancelledState); |
||||||
|
|
||||||
|
emrAddStepsTask.handle(); |
||||||
|
Assert.assertEquals(EXIT_CODE_KILL, emrAddStepsTask.getExitStatusCode()); |
||||||
|
} |
||||||
|
|
||||||
|
@Test |
||||||
|
public void testHandleError() throws Exception { |
||||||
|
when(step.getStatus()).thenReturn(pendingState, runningState, failedState); |
||||||
|
emrAddStepsTask.handle(); |
||||||
|
Assert.assertEquals(EXIT_CODE_FAILURE, emrAddStepsTask.getExitStatusCode()); |
||||||
|
|
||||||
|
when(emrClient.addJobFlowSteps(any())).thenThrow(new AmazonElasticMapReduceException("error"), new EmrTaskException()); |
||||||
|
emrAddStepsTask.handle(); |
||||||
|
Assert.assertEquals(EXIT_CODE_FAILURE, emrAddStepsTask.getExitStatusCode()); |
||||||
|
} |
||||||
|
|
||||||
|
private EmrParameters buildEmrTaskParameters() { |
||||||
|
EmrParameters emrParameters = new EmrParameters(); |
||||||
|
String stepsDefineJson; |
||||||
|
try (InputStream i = this.getClass().getResourceAsStream("EmrAddStepsDefine.json")) { |
||||||
|
assert i != null; |
||||||
|
stepsDefineJson = IOUtils.toString(i, StandardCharsets.UTF_8); |
||||||
|
} catch (Exception e) { |
||||||
|
throw new RuntimeException(e); |
||||||
|
} |
||||||
|
emrParameters.setProgramType(ProgramType.ADD_JOB_FLOW_STEPS); |
||||||
|
emrParameters.setStepsDefineJson(stepsDefineJson); |
||||||
|
|
||||||
|
return emrParameters; |
||||||
|
} |
||||||
|
|
||||||
|
private EmrParameters buildErrorEmrTaskParameters() { |
||||||
|
EmrParameters emrParameters = new EmrParameters(); |
||||||
|
String stepsDefineJson; |
||||||
|
try (InputStream i = this.getClass().getResourceAsStream("EmrErrorAddStepsDefine.json")) { |
||||||
|
assert i != null; |
||||||
|
stepsDefineJson = IOUtils.toString(i, StandardCharsets.UTF_8); |
||||||
|
} catch (Exception e) { |
||||||
|
throw new RuntimeException(e); |
||||||
|
} |
||||||
|
emrParameters.setProgramType(ProgramType.ADD_JOB_FLOW_STEPS); |
||||||
|
emrParameters.setStepsDefineJson(stepsDefineJson); |
||||||
|
|
||||||
|
return emrParameters; |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,17 @@ |
|||||||
|
{ |
||||||
|
"JobFlowId": "j-3V628TKAERHP8", |
||||||
|
"Steps": [ |
||||||
|
{ |
||||||
|
"Name": "calculate_pi", |
||||||
|
"ActionOnFailure": "CONTINUE", |
||||||
|
"HadoopJarStep": { |
||||||
|
"Jar": "command-runner.jar", |
||||||
|
"Args": [ |
||||||
|
"/usr/lib/spark/bin/run-example", |
||||||
|
"SparkPi", |
||||||
|
"15" |
||||||
|
] |
||||||
|
} |
||||||
|
} |
||||||
|
] |
||||||
|
} |
@ -0,0 +1,29 @@ |
|||||||
|
{ |
||||||
|
"JobFlowId": "j-3V628TKAERHP8", |
||||||
|
"Steps": [ |
||||||
|
{ |
||||||
|
"Name": "calculate_pi", |
||||||
|
"ActionOnFailure": "CONTINUE", |
||||||
|
"HadoopJarStep": { |
||||||
|
"Jar": "command-runner.jar", |
||||||
|
"Args": [ |
||||||
|
"/usr/lib/spark/bin/run-example", |
||||||
|
"SparkPi", |
||||||
|
"15" |
||||||
|
] |
||||||
|
} |
||||||
|
}, |
||||||
|
{ |
||||||
|
"Name": "calculate_pi", |
||||||
|
"ActionOnFailure": "CONTINUE", |
||||||
|
"HadoopJarStep": { |
||||||
|
"Jar": "command-runner.jar", |
||||||
|
"Args": [ |
||||||
|
"/usr/lib/spark/bin/run-example", |
||||||
|
"SparkPi", |
||||||
|
"15" |
||||||
|
] |
||||||
|
} |
||||||
|
} |
||||||
|
] |
||||||
|
} |
Loading…
Reference in new issue