alluxio-master-2.x/job/server/src/test/java/alluxio/job/command/CommandHandlingExecutorTest.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.job.command;

import static org.mockito.Mockito.any;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;

import alluxio.AlluxioMockUtil;
import alluxio.client.file.FileSystem;
import alluxio.client.file.FileSystemContext;
import alluxio.grpc.JobCommand;
import alluxio.grpc.RunTaskCommand;
import alluxio.job.JobConfig;
import alluxio.job.JobServerContext;
import alluxio.job.RunTaskContext;
import alluxio.job.TestPlanConfig;
import alluxio.job.util.SerializationUtils;
import alluxio.job.wire.JobWorkerHealth;
import alluxio.underfs.UfsManager;
import alluxio.wire.WorkerNetAddress;
import alluxio.worker.job.JobMasterClient;
import alluxio.worker.job.command.CommandHandlingExecutor;
import alluxio.worker.job.task.TaskExecutorManager;

import com.google.common.collect.Lists;
import com.google.protobuf.ByteString;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import java.io.Serializable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Tests {@link CommandHandlingExecutor}.
 */
public final class CommandHandlingExecutorTest {
  private CommandHandlingExecutor mCommandHandlingExecutor;
  private JobMasterClient mJobMasterClient;
  private TaskExecutorManager mTaskExecutorManager;
  private UfsManager mUfsManager;
  private FileSystemContext mFileSystemContext;
  private FileSystem mFileSystem;

  @Before
  public void before() {

    mJobMasterClient = mock(JobMasterClient.class);
    mTaskExecutorManager = mock(TaskExecutorManager.class);
    WorkerNetAddress workerNetAddress = mock(WorkerNetAddress.class);
    mUfsManager = mock(UfsManager.class);
    mFileSystemContext = mock(FileSystemContext.class);
    mFileSystem = mock(FileSystem.class);
    JobServerContext ctx = new JobServerContext(mFileSystem, mFileSystemContext, mUfsManager);
    mCommandHandlingExecutor =
        new CommandHandlingExecutor(ctx, mTaskExecutorManager, mJobMasterClient, workerNetAddress);
  }

  @Test
  public void heartbeat() throws Exception {
    JobCommand.Builder command = JobCommand.newBuilder();
    RunTaskCommand.Builder runTaskCommand = RunTaskCommand.newBuilder();
    long jobId = 1;
    runTaskCommand.setJobId(jobId);
    long taskId = 2;
    runTaskCommand.setTaskId(taskId);
    JobConfig jobConfig = new TestPlanConfig("/test");
    runTaskCommand.setJobConfig(ByteString.copyFrom(SerializationUtils.serialize(jobConfig)));
    Serializable taskArgs = Lists.newArrayList(1);
    runTaskCommand.setTaskArgs(ByteString.copyFrom(SerializationUtils.serialize(taskArgs)));

    command.setRunTaskCommand(runTaskCommand);

    Mockito.when(mJobMasterClient.heartbeat(any(JobWorkerHealth.class), eq(Lists.newArrayList())))
        .thenReturn(Lists.newArrayList(command.build()));

    mCommandHandlingExecutor.heartbeat(Long.MAX_VALUE);
    ExecutorService executorService = AlluxioMockUtil.getInternalState(
        mCommandHandlingExecutor, "mCommandHandlingService");
    executorService.shutdown();
    Assert.assertTrue(executorService.awaitTermination(5000, TimeUnit.MILLISECONDS));

    Mockito.verify(mTaskExecutorManager).getAndClearTaskUpdates();
    Mockito.verify(mTaskExecutorManager).executeTask(Mockito.eq(jobId), Mockito.eq(taskId),
        Mockito.eq(runTaskCommand.build()), Mockito.any(RunTaskContext.class));
  }
}
