alluxio-master-2.x/job/server/src/test/java/alluxio/master/job/tracker/CliRunnerTest.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.master.job.tracker;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.when;

import alluxio.AlluxioURI;
import alluxio.client.WriteType;
import alluxio.client.file.FileSystem;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.URIStatus;
import alluxio.collections.Pair;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.ListStatusPOptions;
import alluxio.grpc.OperationType;
import alluxio.job.plan.load.LoadConfig;
import alluxio.job.plan.migrate.MigrateConfig;
import alluxio.job.wire.JobSource;
import alluxio.master.job.JobMaster;
import alluxio.master.job.common.CmdInfo;
import alluxio.master.job.metrics.DistributedCmdMetrics;
import alluxio.retry.CountingRetry;
import alluxio.util.CommonUtils;

import com.google.common.base.Objects;
import com.google.common.collect.Lists;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.mockito.MockedStatic;
import org.mockito.Mockito;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * Tests {@link CliRunnerTest}.
 */
public final class CliRunnerTest {
  private static final long DEFAULT_FILE_SIZE = 100;

  private DistLoadCliRunner mLoadRunner;
  private MigrateCliRunner mMigrateRunner;
  private JobMaster mJobMaster;
  private FileSystem mFs;
  private MockedStatic<FileSystem.Factory> mMockStaticFactory;

  @Rule
  public ExpectedException mException = ExpectedException.none();

  @Before
  public void before() throws Exception {
    mFs = mock(FileSystem.class);
    FileSystemContext fsCtx = mock(FileSystemContext.class);
    mMockStaticFactory = mockStatic(FileSystem.Factory.class);
    mMockStaticFactory.when(() ->
        FileSystem.Factory.create(any(FileSystemContext.class))).thenReturn(mFs);

    AlluxioConfiguration conf = mock(AlluxioConfiguration.class);

    when(fsCtx.getPathConf(any(AlluxioURI.class))).thenReturn(conf);
    when(conf.getEnum(any(PropertyKey.class), any()))
            .thenReturn(WriteType.THROUGH);

    mJobMaster = mock(JobMaster.class);
    mLoadRunner = new DistLoadCliRunner(fsCtx, mJobMaster);
    mMigrateRunner = new MigrateCliRunner(fsCtx, mJobMaster);
  }

  @After
  public void after() {
    mMockStaticFactory.close();
  }

  @Test
  public void testRunDistLoad() throws Exception {
    long fileLength = 10;
    int fileCount = 1;

    String srcString = "/src";
    AlluxioURI src = new AlluxioURI(srcString);

    URIStatus srcStatus = mock(URIStatus.class);
    List<URIStatus> listedUri = Lists.newArrayList();
    listedUri.add(srcStatus);
    List<String> listedPath = Lists.newArrayList();
    listedPath.add(srcString);

    int batchSize = 1;
    int replication = 1;
    Set<String> workerSet = Collections.EMPTY_SET;
    Set<String> excludedWorkerSet = Collections.EMPTY_SET;
    Set<String> localityIds = Collections.EMPTY_SET;
    Set<String> excludedLocalityIds = Collections.EMPTY_SET;
    boolean directCache = false;
    long jobControlId = 100;

    when(mFs.getStatus(src)).thenReturn(srcStatus);
    when(srcStatus.getLength()).thenReturn(fileLength);
    when(srcStatus.isFolder()).thenReturn(false);
    when(srcStatus.isCompleted()).thenReturn(true);
    when(srcStatus.getPath()).thenReturn(srcString);
    when(srcStatus.getInAlluxioPercentage()).thenReturn(0);

    //mock the callback on FileSystem.iterateStatus
    Mockito.doAnswer(ans -> {
      Consumer<URIStatus> callback = ans.getArgument(2);
      callback.accept(srcStatus);
      return null;
    }).when(mFs).iterateStatus(any(AlluxioURI.class),
            any(ListStatusPOptions.class), any(Consumer.class));

    // The actual CmdInfo after running through MigrateRunner.
    CmdInfo cmdInfo = mLoadRunner.runDistLoad(batchSize, src, replication, workerSet,
            excludedWorkerSet, localityIds,
            excludedLocalityIds, directCache, jobControlId);
    //only 1 attempt in this test
    CmdRunAttempt actualSingleAttempt = cmdInfo.getCmdRunAttempt().get(0);

    CmdRunAttempt expectedAttempt = new CmdRunAttempt(new CountingRetry(3), mJobMaster);
    expectedAttempt.setFileCount(fileCount);
    expectedAttempt.setFileSize(fileLength);
    LoadConfig config = new LoadConfig(srcString, replication, workerSet,
            excludedWorkerSet, localityIds, excludedLocalityIds, directCache);
    expectedAttempt.setConfig(config);

    // The expected CmdInfo
    CmdInfo expected = new CmdInfo(jobControlId, OperationType.DIST_LOAD,
            JobSource.CLI, 0, listedPath);
    expected.addCmdRunAttempt(expectedAttempt);

    boolean attemptComparison = compareCmdRunAttempt(expectedAttempt, actualSingleAttempt);
    boolean compareCmdMetaInfo = compareCmdMetaInfo(cmdInfo, expected);
    Assert.assertTrue(attemptComparison);
    Assert.assertTrue(compareCmdMetaInfo);
  }

  @Test
  public void testDistLoadSetJobConfigAndFileMetrics() throws Exception {
    int fileCountLimit = 10;
    int fileNameLength = 5;
    List<URIStatus> filePath = createURIStatuses(fileCountLimit, fileNameLength);
    long expectedFileSize = filePath.size() * DEFAULT_FILE_SIZE;
    String expectedFilePath = filePath.stream()
            .map(URIStatus::getPath).collect(Collectors.joining(","));

    CmdRunAttempt attempt = new CmdRunAttempt(new CountingRetry(3), mJobMaster);
    try (MockedStatic<DistributedCmdMetrics> mockedDistributedCmdMetrics =
             mockStatic(DistributedCmdMetrics.class)) {
      mockedDistributedCmdMetrics.when(() ->
          DistributedCmdMetrics.getFileSize(anyString(), any(FileSystem.class),
              any(CountingRetry.class))).thenReturn(DEFAULT_FILE_SIZE);
      mLoadRunner.setJobConfigAndFileMetrics(filePath, 1, Collections.EMPTY_SET,
          Collections.EMPTY_SET, Collections.EMPTY_SET, Collections.EMPTY_SET, false, attempt);
      Assert.assertEquals(attempt.getFileSize(), expectedFileSize);
      Assert.assertEquals(attempt.getFilePath(), expectedFilePath);
    }
  }

  @Test
  public void testRunDistCp() throws Exception {
    long fileLength = 10;
    int fileCount = 1;

    String srcString = "/src";
    String dstString = "/dst";
    AlluxioURI src = new AlluxioURI(srcString);
    AlluxioURI dst = new AlluxioURI(dstString);

    URIStatus srcStatus = mock(URIStatus.class);
    List<URIStatus> listedUri = Lists.newArrayList();
    listedUri.add(srcStatus);
    List<String> listedPath = Lists.newArrayList();
    listedPath.add(srcString);

    boolean overwrite = false;
    int batchSize = 1;
    long jobControlId = 100;

    when(mFs.getStatus(src)).thenReturn(srcStatus);
    when(srcStatus.isFolder()).thenReturn(false);
    when(mFs.listStatus(src)).thenReturn(listedUri);
    when(srcStatus.getPath()).thenReturn(srcString);
    when(srcStatus.getLength()).thenReturn(fileLength);

    // The actual CmdInfo after running through MigrateRunner.
    CmdInfo cmdInfo = mMigrateRunner.runDistCp(src, dst, overwrite, batchSize, jobControlId);
    //only 1 attempt in this test
    CmdRunAttempt actualSingleAttempt = cmdInfo.getCmdRunAttempt().get(0);

    CmdRunAttempt expectedAttempt = new CmdRunAttempt(new CountingRetry(3), mJobMaster);
    expectedAttempt.setFileCount(fileCount);
    expectedAttempt.setFileSize(fileLength);
    MigrateConfig config = new MigrateConfig(srcString, dstString, WriteType.THROUGH, overwrite);
    expectedAttempt.setConfig(config);

    // The expected CmdInfo
    CmdInfo expected = new CmdInfo(jobControlId, OperationType.DIST_CP,
            JobSource.CLI, 0, listedPath);
    expected.addCmdRunAttempt(expectedAttempt);

    boolean attemptComparison = compareCmdRunAttempt(expectedAttempt, actualSingleAttempt);
    boolean compareCmdMetaInfo = compareCmdMetaInfo(cmdInfo, expected);

    Assert.assertTrue(attemptComparison);
    Assert.assertTrue(compareCmdMetaInfo);
  }

  @Test
  public void testDistCpSetJobConfigAndFileMetrics() {
    int fileCountLimit = 10;
    int fileNameLength = 5;
    List<Pair<String, String>> filePath = createFilePaths(fileCountLimit, fileNameLength);
    long expectedFileSize = filePath.size() * DEFAULT_FILE_SIZE;
    String expectedFilePath = filePath.stream()
        .map(Pair::getFirst).collect(Collectors.joining(","));

    CmdRunAttempt attempt = new CmdRunAttempt(new CountingRetry(3), mJobMaster);
    try (MockedStatic<DistributedCmdMetrics> mockedDistributedCmdMetrics =
             mockStatic(DistributedCmdMetrics.class)) {
      mockedDistributedCmdMetrics.when(() ->
          DistributedCmdMetrics.getFileSize(anyString(), any(FileSystem.class),
              any(CountingRetry.class))).thenReturn(DEFAULT_FILE_SIZE);
      mMigrateRunner.setJobConfigAndFileMetrics(filePath, false, WriteType.THROUGH, attempt);
      Assert.assertEquals(attempt.getFileSize(), expectedFileSize);
      Assert.assertEquals(attempt.getFilePath(), expectedFilePath);
    }
  }

  private List<Pair<String, String>> createFilePaths(int fileCountLimit, int nameLength) {
    List<Pair<String, String>> parts = new ArrayList<>();
    for (int i = 0; i < ThreadLocalRandom.current().nextInt(1, fileCountLimit); i++) {
      String src = CommonUtils.randomAlphaNumString(nameLength);
      String dst = CommonUtils.randomAlphaNumString(nameLength);
      parts.add(new Pair<>(src, dst));
    }
    return parts;
  }

  private List<URIStatus> createURIStatuses(int fileCountLimit, int nameLength) throws Exception {
    List<URIStatus> parts = new ArrayList<>();
    for (int i = 0; i < ThreadLocalRandom.current().nextInt(1, fileCountLimit); i++) {
      String src = CommonUtils.randomAlphaNumString(nameLength);
      URIStatus mock = mock(URIStatus.class);
      when(mock.getPath()).thenReturn(src);
      parts.add(mock);
    }
    return parts;
  }

  private boolean compareCmdRunAttempt(CmdRunAttempt attempt, CmdRunAttempt other) {
    return Objects.equal(attempt.getFileSize(), other.getFileSize())
            && Objects.equal(attempt.getFileCount(), other.getFileCount())
            && Objects.equal(attempt.getJobConfig(), other.getJobConfig());
  }

  private boolean compareCmdMetaInfo(CmdInfo cmdInfo, CmdInfo other) {
    return Objects.equal(cmdInfo.getJobControlId(), other.getJobControlId())
            && Objects.equal(cmdInfo.getJobSource(), other.getJobSource())
            && Objects.equal(cmdInfo.getOperationType(), other.getOperationType());
  }
}
