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

import static org.junit.Assert.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import alluxio.AlluxioURI;
import alluxio.client.file.URIStatus;
import alluxio.collections.Pair;
import alluxio.job.JobServerContext;
import alluxio.job.SelectExecutorsContext;
import alluxio.job.plan.SelectExecutorsTest;
import alluxio.wire.WorkerInfo;

import org.apache.commons.io.FileUtils;
import org.junit.Test;

import java.util.ArrayList;
import java.util.List;
import java.util.PrimitiveIterator;
import java.util.Random;
import java.util.Set;

public class CompactDefinitionSelectExecutorsTest extends SelectExecutorsTest {

  private static final String INPUT_DIR = "/input";
  private static final String OUTPUT_DIR = "/output";
  private static final PrimitiveIterator.OfLong LONG_STREAM
      = (new Random()).longs(FileUtils.ONE_GB * 2, FileUtils.ONE_GB * 10).iterator();

  @Test
  public void testExecutorsParallel() throws Exception {
    int tasksPerWorker = 10;
    int numCompactedFiles = 100;
    int totalFiles = 5000;

    PartitionInfo mockPartitionInfo = mock(PartitionInfo.class);
    when(mockPartitionInfo.getFormat(any())).thenReturn(Format.CSV);

    CompactConfig config = new CompactConfig(mockPartitionInfo, INPUT_DIR, mockPartitionInfo,
        OUTPUT_DIR, numCompactedFiles, 2 * FileUtils.ONE_GB);

    List<URIStatus> inputFiles = new ArrayList<>();
    for (int i = 0; i < totalFiles; i++) {
      inputFiles.add(newFile(Integer.toString(i)));
    }

    when(mMockFileSystem.listStatus(new AlluxioURI(INPUT_DIR))).thenReturn(inputFiles);

    Set<Pair<WorkerInfo, ArrayList<CompactTask>>> result = new CompactDefinition().selectExecutors(
        config, SelectExecutorsTest.JOB_WORKERS, new SelectExecutorsContext(1,
            new JobServerContext(mMockFileSystem, mMockFileSystemContext, mMockUfsManager)));
    assertEquals(JOB_WORKERS.size() * tasksPerWorker, result.size());

    int allCompactTasks = 0;
    for (Pair<WorkerInfo, ArrayList<CompactTask>> tasks : result) {
      allCompactTasks += tasks.getSecond().size();
    }
    assertEquals(numCompactedFiles, allCompactTasks);
  }

  private URIStatus newFile(String name) {
    URIStatus mockFileStatus = mock(URIStatus.class);
    when(mockFileStatus.isFolder()).thenReturn(false);
    when(mockFileStatus.getName()).thenReturn(name);
    when(mockFileStatus.getLength()).thenReturn(LONG_STREAM.next());
    return mockFileStatus;
  }
}
