BatchRequestJob extends AbstractLinearExecutionJob { @Override protected void doWork(Map<String, Object> properties) throws AmbariException { Long executionId = properties.get(BATCH_REQUEST_EXECUTION_ID_KEY) != null ? (Long) properties.get(BATCH_REQUEST_EXECUTION_ID_KEY) : null; Long batchId = properties.get(BATCH_REQUEST_BATCH_ID_KEY) != null ? (Long) properties.get(BATCH_REQUEST_BATCH_ID_KEY) : null; String clusterName = (String) properties.get(BATCH_REQUEST_CLUSTER_NAME_KEY); if (executionId == null || batchId == null) { throw new AmbariException("Unable to retrieve persisted batch request" + ", execution_id = " + executionId + ", batch_id = " + batchId); } Map<String, Integer> taskCounts = getTaskCountProperties(properties); Long requestId = executionScheduleManager.executeBatchRequest (executionId, batchId, clusterName); if (requestId != null) { HostRoleStatus status; BatchRequestResponse batchRequestResponse; do { batchRequestResponse = executionScheduleManager .getBatchRequestResponse(requestId, clusterName); status = HostRoleStatus.valueOf(batchRequestResponse.getStatus()); executionScheduleManager.updateBatchRequest(executionId, batchId, clusterName, batchRequestResponse, true); try { Thread.sleep(statusCheckInterval); } catch (InterruptedException e) { String message = "Job Thread interrupted"; LOG.error(message, e); throw new AmbariException(message, e); } } while (!status.isCompletedState()); Map<String, Integer> aggregateCounts = addTaskCountToProperties (properties, taskCounts, batchRequestResponse); if (executionScheduleManager.hasToleranceThresholdExceeded (executionId, clusterName, aggregateCounts)) { throw new AmbariException("Task failure tolerance limit exceeded" + ", execution_id = " + executionId + ", processed batch_id = " + batchId + ", failed tasks = " + aggregateCounts.get(BATCH_REQUEST_FAILED_TASKS_KEY) + ", total tasks completed = " + aggregateCounts.get(BATCH_REQUEST_TOTAL_TASKS_KEY)); } } } @Inject  BatchRequestJob(ExecutionScheduleManager executionScheduleManager,
                         @Named("statusCheckInterval") long statusCheckInterval);  static final String BATCH_REQUEST_EXECUTION_ID_KEY; static final String BATCH_REQUEST_BATCH_ID_KEY; static final String BATCH_REQUEST_CLUSTER_NAME_KEY; static final String BATCH_REQUEST_FAILED_TASKS_KEY; static final String BATCH_REQUEST_TOTAL_TASKS_KEY; }
@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 = EasyMock.newCapture(); Capture<Long> batchIdCapture = EasyMock.newCapture(); Capture<String> clusterNameCapture = EasyMock.newCapture(); 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()); }