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

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.when;

import alluxio.client.block.BlockStoreClient;
import alluxio.client.file.FileSystem;
import alluxio.client.file.FileSystemContext;
import alluxio.collections.Pair;
import alluxio.job.JobServerContext;
import alluxio.job.SelectExecutorsContext;
import alluxio.job.util.SerializableVoid;
import alluxio.underfs.UfsManager;
import alluxio.wire.BlockInfo;
import alluxio.wire.BlockLocation;
import alluxio.wire.WorkerInfo;
import alluxio.wire.WorkerNetAddress;

import com.beust.jcommander.internal.Sets;
import com.google.common.collect.Lists;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.MockedStatic;

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

/**
 * Tests evict functionality of {@link SetReplicaDefinition}.
 */
public final class SetReplicaDefinitionEvictTest {
  private static final long TEST_BLOCK_ID = 1L;
  private static final WorkerNetAddress ADDRESS_1 =
      new WorkerNetAddress().setHost("host1").setDataPort(10);
  private static final WorkerNetAddress ADDRESS_2 =
      new WorkerNetAddress().setHost("host2").setDataPort(10);
  private static final WorkerNetAddress ADDRESS_3 =
      new WorkerNetAddress().setHost("host3").setDataPort(10);
  private static final WorkerInfo WORKER_INFO_1 = new WorkerInfo().setAddress(ADDRESS_1);
  private static final WorkerInfo WORKER_INFO_2 = new WorkerInfo().setAddress(ADDRESS_2);
  private static final WorkerInfo WORKER_INFO_3 = new WorkerInfo().setAddress(ADDRESS_3);
  private static final Set<Pair<WorkerInfo, SerializableVoid>> EMPTY = Sets.newHashSet();

  private FileSystem mMockFileSystem;
  private FileSystemContext mMockFileSystemContext;
  private BlockStoreClient mMockBlockStore;
  private JobServerContext mJobServerContext;
  private MockedStatic<BlockStoreClient> mMockStaticBlockStore;

  @Before
  public void before() {
    mMockFileSystemContext = mock(FileSystemContext.class);
    mMockFileSystem = mock(FileSystem.class);
    mMockBlockStore = mock(BlockStoreClient.class);
    mJobServerContext = new JobServerContext(mMockFileSystem, mMockFileSystemContext,
        mock(UfsManager.class));
    mMockStaticBlockStore = mockStatic(BlockStoreClient.class);
    mMockStaticBlockStore.when(() -> BlockStoreClient.create(mMockFileSystemContext))
        .thenReturn(mMockBlockStore);
  }

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

  /**
   * Helper function to select executors.
   *
   * @param blockLocations where the block is stored currently
   * @param replicas how many replicas to evict
   * @param workerInfoList a list of currently available job workers
   * @return the selection result
   */
  private Set<Pair<WorkerInfo, SetReplicaTask>> selectExecutorsTestHelper(
      List<BlockLocation> blockLocations, int replicas, List<WorkerInfo> workerInfoList)
      throws Exception {
    BlockInfo blockInfo = new BlockInfo().setBlockId(TEST_BLOCK_ID);
    blockInfo.setLocations(blockLocations);
    when(mMockBlockStore.getInfo(TEST_BLOCK_ID)).thenReturn(blockInfo);

    SetReplicaConfig config = new SetReplicaConfig("", TEST_BLOCK_ID, replicas);
    SetReplicaDefinition definition = new SetReplicaDefinition();
    return definition.selectExecutors(config, workerInfoList,
        new SelectExecutorsContext(1, mJobServerContext));
  }

  @Test
  public void selectExecutorsNoBlockWorkerHasBlock() throws Exception {
    Set<Pair<WorkerInfo, SetReplicaTask>> result =
        selectExecutorsTestHelper(Lists.<BlockLocation>newArrayList(), 0,
            Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    // Expect none as no worker has this copy
    Assert.assertEquals(EMPTY, result);
  }

  @Test
  public void selectExecutorsNoJobWorkerHasBlock() throws Exception {
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)), 0,
        Lists.newArrayList(WORKER_INFO_2, WORKER_INFO_3));
    // Expect none as no worker that is available has this copy
    Assert.assertEquals(EMPTY, result);
  }

  @Test
  public void selectExecutorsOnlyOneBlockWorkerHasBlock() throws Exception {
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)), 0,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = Sets.newHashSet();
    expected.add(new Pair<>(WORKER_INFO_1, new SetReplicaTask(Mode.EVICT)));
    // Expect the only worker 1 having this block
    Assert.assertEquals(expected, result);
  }

  @Test
  public void selectExecutorsAnyOneWorkers() throws Exception {
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(Lists
            .newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1),
                new BlockLocation().setWorkerAddress(ADDRESS_2),
                new BlockLocation().setWorkerAddress(ADDRESS_3)), 2,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    // Expect one worker from all workers having this block
    Assert.assertEquals(1, result.size());
    Assert.assertEquals(new SetReplicaTask(Mode.EVICT), result.iterator().next().getSecond());
  }

  @Test
  public void selectExecutorsAllWorkers() throws Exception {
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(Lists
            .newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1),
                new BlockLocation().setWorkerAddress(ADDRESS_2),
                new BlockLocation().setWorkerAddress(ADDRESS_3)), 0,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = Sets.newHashSet();
    expected.add(new Pair<>(WORKER_INFO_1, new SetReplicaTask(Mode.EVICT)));
    expected.add(new Pair<>(WORKER_INFO_2, new SetReplicaTask(Mode.EVICT)));
    expected.add(new Pair<>(WORKER_INFO_3, new SetReplicaTask(Mode.EVICT)));
    // Expect all workers are selected as they all have this block
    Assert.assertEquals(expected, result);
  }

  @Test
  public void selectExecutorsBothWorkers() throws Exception {
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(Lists
            .newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1),
                new BlockLocation().setWorkerAddress(ADDRESS_2)), 0,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    Set<Pair<WorkerInfo, SetReplicaTask>> expected = Sets.newHashSet();
    expected.add(new Pair<>(WORKER_INFO_1, new SetReplicaTask(Mode.EVICT)));
    expected.add(new Pair<>(WORKER_INFO_2, new SetReplicaTask(Mode.EVICT)));
    // Expect both workers having this block should be selected
    Assert.assertEquals(expected, result);
  }

  @Test
  public void selectExecutorsTargetEqualNumBlocks() throws Exception {
    Set<Pair<WorkerInfo, SetReplicaTask>> result = selectExecutorsTestHelper(
        Lists.newArrayList(new BlockLocation().setWorkerAddress(ADDRESS_1)), 1,
        Lists.newArrayList(WORKER_INFO_1, WORKER_INFO_2, WORKER_INFO_3));
    Assert.assertEquals(EMPTY, result);
  }
}
