content: /*
   * 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 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;
  
  import java.util.HashMap;
  import java.util.Map;
  
  import org.apache.ambari.server.actionmanager.HostRoleStatus;
  import org.apache.ambari.server.scheduler.ExecutionScheduleManager;
  import org.easymock.Capture;
  import org.easymock.EasyMock;
  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;
  
  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<>();
      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_FAILED_TASKS_IN_CURRENT_BATCH_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());
    }
  
    @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<>();
      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 = EasyMock.newCapture();
      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));
    }
  }
  