/*
 * 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.migrate;

import static org.mockito.Mockito.when;

import alluxio.AlluxioURI;
import alluxio.client.WriteType;
import alluxio.client.file.URIStatus;
import alluxio.collections.Pair;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.FileDoesNotExistException;
import alluxio.job.JobServerContext;
import alluxio.job.SelectExecutorsContext;
import alluxio.job.plan.SelectExecutorsTest;
import alluxio.membership.WorkerClusterView;
import alluxio.wire.BlockInfo;
import alluxio.wire.BlockLocation;
import alluxio.wire.FileBlockInfo;
import alluxio.wire.FileInfo;
import alluxio.wire.WorkerInfo;
import alluxio.wire.WorkerNetAddress;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;
import java.util.Set;

/**
 * Unit tests for
 * {@link MigrateDefinition#selectExecutors(MigrateConfig, List, SelectExecutorsContext)}.
 * No matter whether to delete source, selectExecutors should have the same behavior.
 */
public final class MigrateDefinitionSelectExecutorsTest extends SelectExecutorsTest {
  private static final WorkerClusterView BLOCK_WORKERS = new WorkerClusterView(
      Arrays.asList(JOB_WORKER_0, JOB_WORKER_1, JOB_WORKER_2, JOB_WORKER_3));

  @Before
  @Override
  public void before() throws Exception {
    super.before();
    when(mMockFileSystemContext.getCachedWorkers()).thenReturn(BLOCK_WORKERS);
    createDirectory("/");
  }

  @Test
  public void migrateToSelf() throws Exception {
    createDirectory("/src");
    Assert.assertEquals(ImmutableSet.of(), assignMigrates("/src", "/src"));
  }

  @Test
  public void assignToLocalWorker() throws Exception {
    createFileWithBlocksOnWorkers("/src", 0);
    setPathToNotExist("/dst");
    Set<Pair<WorkerInfo, MigrateCommand>> expected = ImmutableSet.of(new Pair<>(
        JOB_WORKERS.get(0), new MigrateCommand("/src", "/dst")));
    Assert.assertEquals(expected, assignMigrates("/src", "/dst"));
  }

  @Test
  public void assignToWorkerWithMostBlocks() throws Exception {
    createFileWithBlocksOnWorkers("/src", 3, 1, 1, 3, 1);
    setPathToNotExist("/dst");
    Set<Pair<WorkerInfo, MigrateCommand>> expected = ImmutableSet.of(new Pair<>(
        JOB_WORKERS.get(1), new MigrateCommand("/src", "/dst")));
    Assert.assertEquals(expected, assignMigrates("/src", "/dst"));
  }

  @Test
  public void migrateToSubpath() throws Exception {
    RuntimeException exception = Assert.assertThrows(RuntimeException.class,
        () -> assignMigratesFail("/src", "/src/dst"));
    Assert.assertEquals(
        ExceptionMessage.MIGRATE_CANNOT_BE_TO_SUBDIRECTORY.getMessage("/src", "/src/dst"),
        exception.getMessage());
  }

  @Test
  public void migrateMissingSource() throws Exception {
    setPathToNotExist("/notExist");
    FileDoesNotExistException exception = Assert.assertThrows(FileDoesNotExistException.class,
        () -> assignMigratesFail("/notExist", "/dst"));
    Assert.assertEquals(ExceptionMessage.PATH_DOES_NOT_EXIST.getMessage("/notExist"),
        exception.getMessage());
  }

  @Test
  public void migrateUncachedFile() throws Exception {
    createFileWithBlocksOnWorkers("/src");
    setPathToNotExist("/dst");
    Assert.assertEquals(1, assignMigrates("/src", "/dst").size());
  }

  @Test
  public void migrateNoLocalJobWorker() throws Exception {
    createFileWithBlocksOnWorkers("/src", 0);
    setPathToNotExist("/dst");

    Set<Pair<WorkerInfo, MigrateCommand>> assignments =
        new MigrateDefinition().selectExecutors(
            new MigrateConfig("/src", "/dst", WriteType.THROUGH, true),
            ImmutableList.of(JOB_WORKER_3),
            new SelectExecutorsContext(1,
                new JobServerContext(mMockFileSystem, mMockFileSystemContext, mMockUfsManager)));

    Assert.assertEquals(ImmutableSet.of(new Pair<>(JOB_WORKER_3,
        new MigrateCommand("/src", "/dst"))), assignments);
  }

  /**
   * Runs selectExecutors for the migrate from source to destination.
   */
  private Set<Pair<WorkerInfo, MigrateCommand>> assignMigrates(String source,
      String destination) throws Exception {
    return assignMigrates(new MigrateConfig(source, destination, WriteType.THROUGH, false));
  }

  /**
   * Runs selectExecutors for the migrate from source to destination with the given writeType and
   * overwrite value.
   */
  private Set<Pair<WorkerInfo, MigrateCommand>> assignMigrates(MigrateConfig config)
      throws Exception {
    return new MigrateDefinition().selectExecutors(config,
        JOB_WORKERS, new SelectExecutorsContext(1,
            new JobServerContext(mMockFileSystem, mMockFileSystemContext, mMockUfsManager)));
  }

  /**
   * Runs selectExecutors with the expectation that it will throw an exception.
   */
  private void assignMigratesFail(String source, String destination) throws Exception {
    assignMigratesFail(source, destination, WriteType.THROUGH, false);
  }

  /**
   * Runs selectExecutors with the expectation that it will throw an exception.
   */
  private void assignMigratesFail(String source, String destination, WriteType writeType,
      boolean overwrite) throws Exception {
    Set<Pair<WorkerInfo, MigrateCommand>> assignment =
        assignMigrates(new MigrateConfig(source, destination, writeType, overwrite));
    Assert.fail(
        "Selecting executors should have failed, but it succeeded with assignment " + assignment);
  }

  private FileInfo createFileWithBlocksOnWorkers(String testFile, int... workerInds)
      throws Exception {
    return createFileWithBlocksOnWorkers(testFile, new FileInfo(), workerInds);
  }

  /**
   * Creates a file with the given name and a block on each specified worker. Workers may be
   * repeated to give them multiple blocks.
   *
   * @param testFile the name of the file to create
   * @param fileInfo file info to apply to the created file
   * @param workerInds the workers to put blocks on, specified by their indices
   * @return file info for the created file
   */
  private FileInfo createFileWithBlocksOnWorkers(String testFile, FileInfo fileInfo,
      int... workerInds) throws Exception {
    AlluxioURI uri = new AlluxioURI(testFile);
    List<FileBlockInfo> blockInfos = Lists.newArrayList();
    for (int workerInd : workerInds) {
      WorkerNetAddress address = JOB_WORKERS.get(workerInd).getAddress();
      blockInfos.add(new FileBlockInfo().setBlockInfo(new BlockInfo()
          .setLocations(Lists.newArrayList(new BlockLocation().setWorkerAddress(address)))));
    }
    FileInfo testFileInfo =
        fileInfo.setFolder(false).setPath(testFile).setFileBlockInfos(blockInfos);
    when(mMockFileSystem.listStatus(uri))
        .thenReturn(Lists.newArrayList(new URIStatus(testFileInfo)));
    when(mMockFileSystem.getStatus(uri)).thenReturn(new URIStatus(testFileInfo));
    return testFileInfo;
  }

  /**
   * Creates a directory with the given name.
   *
   * @return file info for the created directory
   */
  private FileInfo createDirectory(String name) throws Exception {
    // Call all directories mount points to force cross-mount functionality.
    FileInfo info = new FileInfo().setFolder(true).setPath(name).setMountPoint(true);
    when(mMockFileSystem.getStatus(new AlluxioURI(name))).thenReturn(new URIStatus(info));
    return info;
  }

  /**
   * Tells the mock to throw FileDoesNotExistException when the given path is queried.
   */
  private void setPathToNotExist(String path) throws Exception {
    AlluxioURI uri = new AlluxioURI(path);
    when(mMockFileSystem.getStatus(uri)).thenThrow(new FileDoesNotExistException(uri));
    when(mMockFileSystem.listStatus(uri)).thenThrow(new FileDoesNotExistException(uri));
  }
}
