/*
* 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.ambari.server.state.scheduler;
import org.apache.ambari.server.actionmanager.HostRoleStatus;
import org.apache.ambari.server.scheduler.ExecutionScheduleManager;
import org.easymock.Capture;
import org.junit.Assert;
import org.junit.Test;
import org.quartz.JobDataMap;
import org.quartz.JobDetail;
import org.quartz.JobExecutionContext;
import org.quartz.JobKey;
import org.quartz.Trigger;
import java.util.HashMap;
import java.util.Map;
import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.capture;
import static org.easymock.EasyMock.captureLong;
import static org.easymock.EasyMock.createMock;
import static org.easymock.EasyMock.createMockBuilder;
import static org.easymock.EasyMock.createNiceMock;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.expectLastCall;
import static org.easymock.EasyMock.replay;
import static org.easymock.EasyMock.verify;
public class BatchRequestJobTest {
@Test
public void testDoWork() throws Exception {
ExecutionScheduleManager scheduleManagerMock = createMock(ExecutionScheduleManager.class);
BatchRequestJob batchRequestJob = new BatchRequestJob(scheduleManagerMock, 100L);
String clusterName = "mycluster";
long requestId = 11L;
long executionId = 31L;
long batchId = 1L;
Map<String, Object> properties = new HashMap<String, Object>();
properties.put(BatchRequestJob.BATCH_REQUEST_EXECUTION_ID_KEY, executionId);
properties.put(BatchRequestJob.BATCH_REQUEST_BATCH_ID_KEY, batchId);
properties.put(BatchRequestJob.BATCH_REQUEST_CLUSTER_NAME_KEY, clusterName);
HashMap<String, Integer> taskCounts = new HashMap<String, Integer>()
{{ put(BatchRequestJob.BATCH_REQUEST_FAILED_TASKS_KEY, 0);
put(BatchRequestJob.BATCH_REQUEST_TOTAL_TASKS_KEY, 0); }};
BatchRequestResponse pendingResponse = new BatchRequestResponse();
pendingResponse.setStatus(HostRoleStatus.PENDING.toString());
BatchRequestResponse inProgressResponse = new BatchRequestResponse();
inProgressResponse.setStatus(HostRoleStatus.IN_PROGRESS.toString());
BatchRequestResponse completedResponse = new BatchRequestResponse();
completedResponse.setStatus(HostRoleStatus.COMPLETED.toString());
Capture<Long> executionIdCapture = new Capture<Long>();
Capture<Long> batchIdCapture = new Capture<Long>();
Capture<String> clusterNameCapture = new Capture<String>();
expect(scheduleManagerMock.executeBatchRequest(captureLong(executionIdCapture),
captureLong(batchIdCapture),
capture(clusterNameCapture))).andReturn(requestId);
expect(scheduleManagerMock.getBatchRequestResponse(requestId, clusterName)).
andReturn(pendingResponse).times(2);
expect(scheduleManagerMock.getBatchRequestResponse(requestId, clusterName)).
andReturn(inProgressResponse).times(4);
expect(scheduleManagerMock.getBatchRequestResponse(requestId, clusterName)).
andReturn(completedResponse).once();
expect(scheduleManagerMock.hasToleranceThresholdExceeded(executionId,
clusterName, taskCounts)).andReturn(false);
scheduleManagerMock.updateBatchRequest(eq(executionId), eq(batchId), eq(clusterName),
anyObject(BatchRequestResponse.class), eq(true));
expectLastCall().anyTimes();
replay(scheduleManagerMock);
batchRequestJob.doWork(properties);
verify(scheduleManagerMock);
Assert.assertEquals(executionId, executionIdCapture.getValue().longValue());
Assert.assertEquals(batchId, batchIdCapture.getValue().longValue());
Assert.assertEquals(clusterName, clusterNameCapture.getValue());
}
@Test
public void testTaskCountsPersistedWithTrigger() throws Exception {
ExecutionScheduleManager scheduleManagerMock = createNiceMock
(ExecutionScheduleManager.class);
BatchRequestJob batchRequestJobMock = createMockBuilder
(BatchRequestJob.class).withConstructor(scheduleManagerMock, 100L)
.addMockedMethods("doWork")
.createMock();
JobExecutionContext executionContext = createNiceMock(JobExecutionContext.class);
JobDataMap jobDataMap = createNiceMock(JobDataMap.class);
JobDetail jobDetail = createNiceMock(JobDetail.class);
Map<String, Object> properties = new HashMap<String, Object>();
properties.put(BatchRequestJob.BATCH_REQUEST_FAILED_TASKS_KEY, 10);
properties.put(BatchRequestJob.BATCH_REQUEST_TOTAL_TASKS_KEY, 20);
expect(scheduleManagerMock.continueOnMisfire(executionContext)).andReturn(true);
expect(executionContext.getMergedJobDataMap()).andReturn(jobDataMap);
expect(executionContext.getJobDetail()).andReturn(jobDetail);
expect(jobDetail.getKey()).andReturn(JobKey.jobKey("testJob", "testGroup"));
expect(jobDataMap.getWrappedMap()).andReturn(properties);
expect(jobDataMap.getString((String) anyObject())).andReturn("testJob").anyTimes();
Capture<Trigger> triggerCapture = new Capture<Trigger>();
scheduleManagerMock.scheduleJob(capture(triggerCapture));
expectLastCall().once();
replay(scheduleManagerMock, executionContext, jobDataMap, jobDetail);
batchRequestJobMock.execute(executionContext);
verify(scheduleManagerMock, executionContext, jobDataMap, jobDetail);
Trigger trigger = triggerCapture.getValue();
Assert.assertNotNull(trigger);
JobDataMap savedMap = trigger.getJobDataMap();
Assert.assertNotNull(savedMap);
Assert.assertEquals(10, savedMap.getIntValue
(BatchRequestJob.BATCH_REQUEST_FAILED_TASKS_KEY));
Assert.assertEquals(20, savedMap.getIntValue
(BatchRequestJob.BATCH_REQUEST_TOTAL_TASKS_KEY));
}
}