content: /*
   * 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.client.block;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertThrows;
  import static org.junit.Assert.assertTrue;
  import static org.mockito.Mockito.any;
  import static org.mockito.Mockito.doAnswer;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  import alluxio.ClientContext;
  import alluxio.ConfigurationRule;
  import alluxio.client.WriteType;
  import alluxio.client.block.policy.BlockLocationPolicy;
  import alluxio.client.block.policy.options.GetWorkerOptions;
  import alluxio.client.block.stream.BlockInStream;
  import alluxio.client.block.stream.BlockOutStream;
  import alluxio.client.block.stream.BlockWorkerClient;
  import alluxio.client.block.stream.BlockWorkerDataReader;
  import alluxio.client.block.stream.GrpcDataReader;
  import alluxio.client.block.stream.NoopClosableResource;
  import alluxio.client.file.FileSystemContext;
  import alluxio.client.file.URIStatus;
  import alluxio.client.file.options.InStreamOptions;
  import alluxio.client.file.options.OutStreamOptions;
  import alluxio.conf.AlluxioConfiguration;
  import alluxio.conf.Configuration;
  import alluxio.conf.InstancedConfiguration;
  import alluxio.conf.PropertyKey;
  import alluxio.exception.ExceptionMessage;
  import alluxio.exception.PreconditionMessage;
  import alluxio.exception.status.UnavailableException;
  import alluxio.grpc.CreateLocalBlockResponse;
  import alluxio.grpc.OpenFilePOptions;
  import alluxio.grpc.OpenLocalBlockRequest;
  import alluxio.grpc.OpenLocalBlockResponse;
  import alluxio.network.TieredIdentityFactory;
  import alluxio.resource.DummyCloseableResource;
  import alluxio.util.FileSystemOptionsUtils;
  import alluxio.util.network.NetworkAddressUtils;
  import alluxio.wire.BlockInfo;
  import alluxio.wire.BlockLocation;
  import alluxio.wire.FileBlockInfo;
  import alluxio.wire.FileInfo;
  import alluxio.wire.WorkerNetAddress;
  import alluxio.worker.block.BlockWorker;
  
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.Lists;
  import com.google.common.collect.Sets;
  import io.grpc.stub.ClientCallStreamObserver;
  import io.grpc.stub.StreamObserver;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.mockito.Mockito;
  import org.mockito.stubbing.Answer;
  import org.powermock.api.mockito.PowerMockito;
  import org.powermock.core.classloader.annotations.PrepareForTest;
  import org.powermock.modules.junit4.PowerMockRunner;
  
  import java.io.Closeable;
  import java.io.File;
  import java.util.AbstractMap;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.stream.Collectors;
  import javax.annotation.concurrent.ThreadSafe;
  
  /**
   * Tests for {@link BlockStoreClient}.
   */
  @RunWith(PowerMockRunner.class)
  @PrepareForTest({FileSystemContext.class})
  public final class BlockStoreClientTest {
  
    private static final InstancedConfiguration S_CONF = Configuration.copyGlobal();
  
    private static final long BLOCK_ID = 3L;
    private static final long BLOCK_LENGTH = 100L;
    private static final String WORKER_HOSTNAME_LOCAL =
        NetworkAddressUtils.getLocalHostName((int) S_CONF
            .getMs(PropertyKey.NETWORK_HOST_RESOLUTION_TIMEOUT_MS));
    private static final String WORKER_HOSTNAME_REMOTE = "remote";
    private static final WorkerNetAddress WORKER_NET_ADDRESS_LOCAL = new WorkerNetAddress()
        .setHost(WORKER_HOSTNAME_LOCAL);
    private static final WorkerNetAddress WORKER_NET_ADDRESS_REMOTE = new WorkerNetAddress()
        .setHost(WORKER_HOSTNAME_REMOTE);
    private ClientCallStreamObserver mStreamObserver;
    private StreamObserver<OpenLocalBlockResponse> mResponseObserver;
  
    /**
     * A mock class used to return controlled result when selecting workers.
     */
    @ThreadSafe
    public static class MockBlockLocationPolicyTest implements BlockLocationPolicy {
      private List<WorkerNetAddress> mWorkerNetAddresses;
      private int mIndex;
  
      /**
       * Cosntructs this mock location policy with empty host list,
       * needed for instantiation in {@link BlockLocationPolicy.Factory}.
       *
       * @param ignoredConf is unused
       */
      public MockBlockLocationPolicyTest(AlluxioConfiguration ignoredConf) {
        mIndex = 0;
        mWorkerNetAddresses =  Collections.emptyList();
      }
  
      /**
       * Constructs this mock policy that returns the given result, once a time, in the input order.
       *
       * @param addresses list of addresses this mock policy will return
       */
      public MockBlockLocationPolicyTest(List<WorkerNetAddress> addresses) {
        mWorkerNetAddresses = Lists.newArrayList(addresses);
        mIndex = 0;
      }
  
      public void setHosts(List<WorkerNetAddress> addresses) {
        mWorkerNetAddresses = Lists.newArrayList(addresses);
        mIndex = 0;
      }
  
      @Override
      public Optional<WorkerNetAddress> getWorker(GetWorkerOptions options) {
        if (mWorkerNetAddresses.isEmpty()) {
          return Optional.empty();
        }
        return Optional.of(mWorkerNetAddresses.get(mIndex++));
      }
    }
  
    private BlockMasterClient mMasterClient;
    private BlockWorkerClient mWorkerClient;
    private BlockStoreClient mBlockStore;
    private FileSystemContext mContext;
    private ClientContext mClientContext;
  
    @Before
    public void before() throws Exception {
      mMasterClient = PowerMockito.mock(BlockMasterClient.class);
      mWorkerClient = PowerMockito.mock(BlockWorkerClient.class);
  
      mClientContext = ClientContext.create(S_CONF);
  
      mContext = PowerMockito.mock(FileSystemContext.class);
      when(mContext.acquireBlockMasterClientResource())
          .thenReturn(new DummyCloseableResource<>(mMasterClient));
      when(mContext.getClientContext()).thenReturn(mClientContext);
      when(mContext.getClusterConf()).thenReturn(S_CONF);
      when(mContext.getReadBlockLocationPolicy(any(AlluxioConfiguration.class)))
          .thenAnswer((Answer) invocation -> {
            AlluxioConfiguration conf =
                invocation.getArgument(0, AlluxioConfiguration.class);
            return BlockLocationPolicy.Factory.create(
                conf.getClass(PropertyKey.USER_UFS_BLOCK_READ_LOCATION_POLICY), conf);
          });
  
      mBlockStore = new BlockStoreClient(mContext,
          TieredIdentityFactory.fromString("node=" + WORKER_HOSTNAME_LOCAL, S_CONF));
  
      when(mContext.acquireBlockWorkerClient(any(WorkerNetAddress.class)))
          .thenReturn(new NoopClosableResource<>(mWorkerClient));
      mStreamObserver = PowerMockito.mock(ClientCallStreamObserver.class);
      when(mWorkerClient.writeBlock(any(StreamObserver.class)))
          .thenReturn(mStreamObserver);
      when(mWorkerClient.openLocalBlock(any(StreamObserver.class)))
          .thenReturn(mStreamObserver);
      when(mStreamObserver.isReady()).thenReturn(true);
      when(mContext.getCachedWorkers()).thenReturn(Lists.newArrayList(
          new BlockWorkerInfo(new WorkerNetAddress(), -1, -1)));
    }
  
    @Test
    public void getOutStreamUsingLocationPolicy() {
      OutStreamOptions options =
          OutStreamOptions.defaults(mContext).setWriteType(WriteType.MUST_CACHE)
              .setLocationPolicy((workerOptions) -> {
                throw new RuntimeException("policy threw exception");
              });
      assertThrows(Exception.class, () -> mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options));
    }
  
    @Test
    public void getOutStreamMissingLocationPolicy() {
      OutStreamOptions options =
          OutStreamOptions.defaults(mContext).setBlockSizeBytes(BLOCK_LENGTH)
              .setWriteType(WriteType.MUST_CACHE).setLocationPolicy(null);
      Exception e = assertThrows(NullPointerException.class, () ->
          mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options));
      assertTrue(e.getMessage()
          .contains(PreconditionMessage.BLOCK_WRITE_LOCATION_POLICY_UNSPECIFIED.toString()));
    }
  
    @Test
    public void getOutStreamNoWorker() {
      OutStreamOptions options =
          OutStreamOptions
              .defaults(mContext)
              .setBlockSizeBytes(BLOCK_LENGTH)
              .setWriteType(WriteType.MUST_CACHE)
              .setLocationPolicy(
                  new MockBlockLocationPolicyTest(Lists.newArrayList()));
      Exception e = assertThrows(UnavailableException.class, () ->
          mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options));
      assertTrue(e.getMessage()
          .contains(ExceptionMessage.NO_SPACE_FOR_BLOCK_ON_WORKER.getMessage(BLOCK_LENGTH)));
    }
  
    @Test
    public void getOutStreamLocal() throws Exception {
      File file = File.createTempFile("test", ".tmp");
      CreateLocalBlockResponse response = CreateLocalBlockResponse.newBuilder()
          .setPath(file.getAbsolutePath()).build();
      when(mWorkerClient.createLocalBlock(any(StreamObserver.class)))
          .thenAnswer((Answer) invocation -> {
            StreamObserver<CreateLocalBlockResponse> observer =
                invocation.getArgument(0, StreamObserver.class);
            observer.onNext(response);
            return mStreamObserver;
          });
  
      OutStreamOptions options = OutStreamOptions.defaults(mContext)
          .setBlockSizeBytes(BLOCK_LENGTH).setLocationPolicy(
              new MockBlockLocationPolicyTest(Lists.newArrayList(WORKER_NET_ADDRESS_LOCAL)))
              .setWriteType(WriteType.MUST_CACHE);
      BlockOutStream stream = mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options);
      assertEquals(WORKER_NET_ADDRESS_LOCAL, stream.getAddress());
    }
  
    @Test
    public void getOutStreamRemote() throws Exception {
      WorkerNetAddress worker1 = new WorkerNetAddress().setHost("worker1");
      WorkerNetAddress worker2 = new WorkerNetAddress().setHost("worker2");
      OutStreamOptions options =
          OutStreamOptions.defaults(mContext).setBlockSizeBytes(BLOCK_LENGTH)
              .setLocationPolicy(new MockBlockLocationPolicyTest(Arrays.asList(worker1, worker2)))
              .setWriteType(WriteType.MUST_CACHE);
      BlockOutStream stream1 = mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options);
      assertEquals(worker1, stream1.getAddress());
      BlockOutStream stream2 = mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options);
      assertEquals(worker2, stream2.getAddress());
    }
  
    @Test
    public void getOutStreamWithReplicated() throws Exception {
      File file = File.createTempFile("test", ".tmp");
      CreateLocalBlockResponse response = CreateLocalBlockResponse.newBuilder()
          .setPath(file.getAbsolutePath()).build();
      when(mWorkerClient.createLocalBlock(any(StreamObserver.class)))
          .thenAnswer((Answer) invocation -> {
            StreamObserver<CreateLocalBlockResponse> observer =
                invocation.getArgument(0, StreamObserver.class);
            observer.onNext(response);
            return mStreamObserver;
          });
  
      when(mContext.getCachedWorkers()).thenReturn(Lists
          .newArrayList(new BlockWorkerInfo(WORKER_NET_ADDRESS_LOCAL, -1, -1),
              new BlockWorkerInfo(WORKER_NET_ADDRESS_REMOTE, -1, -1)));
      OutStreamOptions options =
          OutStreamOptions.defaults(mContext).setBlockSizeBytes(BLOCK_LENGTH).setLocationPolicy(
              new MockBlockLocationPolicyTest(
                  Lists.newArrayList(WORKER_NET_ADDRESS_LOCAL, WORKER_NET_ADDRESS_REMOTE)))
              .setWriteType(WriteType.MUST_CACHE).setReplicationMin(2);
      BlockOutStream stream = mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options);
  
      assertEquals(alluxio.client.block.stream.BlockOutStream.class, stream.getClass());
    }
  
    @Test
    public void getInStreamUfsMockLocaltion() throws Exception {
      try (Closeable ignored = new ConfigurationRule(PropertyKey.USER_UFS_BLOCK_READ_LOCATION_POLICY,
          MockBlockLocationPolicyTest.class.getTypeName(), S_CONF).toResource()) {
        WorkerNetAddress worker1 = new WorkerNetAddress().setHost("worker1");
        WorkerNetAddress worker2 = new WorkerNetAddress().setHost("worker2");
        BlockInfo info = new BlockInfo().setBlockId(0);
        URIStatus dummyStatus = new URIStatus(new FileInfo().setPersisted(true)
            .setBlockIds(Collections.singletonList(0L))
            .setFileBlockInfos(Collections.singletonList(new FileBlockInfo().setBlockInfo(info))));
        OpenFilePOptions readOptions = OpenFilePOptions.newBuilder().build();
        InStreamOptions options = new InStreamOptions(dummyStatus, readOptions, S_CONF, mContext);
        ((MockBlockLocationPolicyTest) options.getUfsReadLocationPolicy())
            .setHosts(Arrays.asList(worker1, worker2));
        when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(new BlockInfo());
        when(mContext.getCachedWorkers()).thenReturn(
            Lists.newArrayList(new BlockWorkerInfo(worker1, -1, -1),
                new BlockWorkerInfo(worker2, -1, -1)));
  
        // Location policy chooses worker1 first.
        assertEquals(worker1, mBlockStore.getInStream(BLOCK_ID, options).getAddress());
        // Location policy chooses worker2 second.
        assertEquals(worker2, mBlockStore.getInStream(BLOCK_ID, options).getAddress());
      }
    }
  
    @Test
    public void getInStreamUfsLocalFirst() throws Exception {
      WorkerNetAddress remote = new WorkerNetAddress().setHost("remote");
      WorkerNetAddress local = new WorkerNetAddress().setHost(WORKER_HOSTNAME_LOCAL);
      BlockInfo info = new BlockInfo().setBlockId(0);
      URIStatus dummyStatus =
          new URIStatus(new FileInfo().setPersisted(true).setBlockIds(Collections.singletonList(0L))
              .setFileBlockInfos(Collections.singletonList(new FileBlockInfo().setBlockInfo(info))));
      OpenFilePOptions readOptions = OpenFilePOptions.newBuilder().build();
      InStreamOptions options = new InStreamOptions(dummyStatus, readOptions, S_CONF, mContext);
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(new BlockInfo());
      when(mContext.getCachedWorkers()).thenReturn(
          Lists.newArrayList(new BlockWorkerInfo(remote, 100, 0),
              new BlockWorkerInfo(local, 100, 0)));
  
      BlockInStream stream = mBlockStore.getInStream(BLOCK_ID, options);
      assertEquals(local, stream.getAddress());
      assertEquals(GrpcDataReader.Factory.class.getName(),
          stream.getDataReaderFactory().getClass().getName());
    }
  
    @Test
    public void getInStreamNoWorkers() throws Exception {
      URIStatus dummyStatus =
          new URIStatus(new FileInfo().setPersisted(true).setBlockIds(Collections.singletonList(0L)));
      InStreamOptions options =
          new InStreamOptions(dummyStatus, FileSystemOptionsUtils.openFileDefaults(S_CONF),
              S_CONF, mContext);
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(new BlockInfo());
      when(mContext.getCachedWorkers()).thenReturn(Collections.emptyList());
      Exception e = assertThrows(UnavailableException.class, () ->
          mBlockStore.getInStream(BLOCK_ID, options).getAddress());
      assertTrue(e.getMessage().contains(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage()));
    }
  
    @Test
    public void getInStreamMissingBlock() throws Exception {
      URIStatus dummyStatus = new URIStatus(
          new FileInfo().setPersisted(false).setBlockIds(Collections.singletonList(0L)));
      InStreamOptions options =
          new InStreamOptions(dummyStatus, FileSystemOptionsUtils.openFileDefaults(S_CONF),
              S_CONF, mContext);
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(new BlockInfo());
      Exception e = assertThrows(UnavailableException.class, () ->
          mBlockStore.getInStream(BLOCK_ID, options).getAddress());
      assertTrue(e.getMessage().contains("unavailable in both Alluxio and UFS"));
    }
  
    @Test
    public void getInStreamLocal() throws Exception {
      WorkerNetAddress remote = new WorkerNetAddress().setHost("remote");
      WorkerNetAddress local = new WorkerNetAddress().setHost(WORKER_HOSTNAME_LOCAL);
  
      // Mock away gRPC usage.
      OpenLocalBlockResponse response = OpenLocalBlockResponse.newBuilder().setPath("/tmp").build();
      when(mWorkerClient.openLocalBlock(any(StreamObserver.class))).thenAnswer(invocation -> {
        mResponseObserver = invocation.getArgument(0, StreamObserver.class);
        return mStreamObserver;
      });
      doAnswer(invocation -> {
        mResponseObserver.onNext(response);
        mResponseObserver.onCompleted();
        return null;
      }).when(mStreamObserver).onNext(any(OpenLocalBlockRequest.class));
  
      BlockInfo info = new BlockInfo().setBlockId(BLOCK_ID).setLocations(Arrays
          .asList(new BlockLocation().setWorkerAddress(remote),
              new BlockLocation().setWorkerAddress(local)));
  
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(info);
      assertEquals(local, mBlockStore.getInStream(BLOCK_ID, new InStreamOptions(
          new URIStatus(new FileInfo().setBlockIds(Lists.newArrayList(BLOCK_ID))),
              S_CONF, mContext))
          .getAddress());
    }
  
    @Test
    public void getInStreamRemote() throws Exception {
      WorkerNetAddress remote1 = new WorkerNetAddress().setHost("remote1");
      WorkerNetAddress remote2 = new WorkerNetAddress().setHost("remote2");
  
      BlockInfo info = new BlockInfo().setBlockId(BLOCK_ID).setLocations(Arrays
          .asList(new BlockLocation().setWorkerAddress(remote1),
              new BlockLocation().setWorkerAddress(remote2)));
  
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(info);
      // We should sometimes get remote1 and sometimes get remote2.
      Set<WorkerNetAddress> results = new HashSet<>();
      for (int i = 0; i < 40; i++) {
        results.add(mBlockStore.getInStream(BLOCK_ID, new InStreamOptions(
            new URIStatus(new FileInfo().setBlockIds(Lists.newArrayList(BLOCK_ID))),
                S_CONF, mContext))
            .getAddress());
      }
      assertEquals(Sets.newHashSet(remote1, remote2), results);
    }
  
    @Test
    public void getInStreamProcessLocal() throws Exception {
      WorkerNetAddress remote = new WorkerNetAddress().setHost("remote");
      WorkerNetAddress local = new WorkerNetAddress().setHost(WORKER_HOSTNAME_LOCAL);
      BlockInfo info = new BlockInfo().setBlockId(BLOCK_ID).setLocations(Arrays
          .asList(new BlockLocation().setWorkerAddress(remote),
              new BlockLocation().setWorkerAddress(local)));
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(info);
  
      when(mContext.hasProcessLocalWorker()).thenReturn(true);
      Optional<BlockWorker> blockWorker = Optional.of(Mockito.mock(BlockWorker.class));
      when(mContext.getProcessLocalWorker()).thenReturn(blockWorker);
  
      BlockInStream stream = mBlockStore.getInStream(BLOCK_ID, new InStreamOptions(
          new URIStatus(new FileInfo().setBlockIds(Lists.newArrayList(BLOCK_ID))),
          S_CONF, mContext));
      assertEquals(local, stream.getAddress());
      assertEquals(BlockWorkerDataReader.Factory.class.getName(),
          stream.getDataReaderFactory().getClass().getName());
    }
  
    @Test
    public void getInStreamUfsProcessLocal() throws Exception {
      WorkerNetAddress remote = new WorkerNetAddress().setHost("remote");
      WorkerNetAddress local = new WorkerNetAddress().setHost(WORKER_HOSTNAME_LOCAL);
      BlockInfo info = new BlockInfo().setBlockId(0);
      URIStatus dummyStatus =
          new URIStatus(new FileInfo().setPersisted(true).setBlockIds(Collections.singletonList(0L))
              .setFileBlockInfos(Collections.singletonList(new FileBlockInfo().setBlockInfo(info))));
      OpenFilePOptions readOptions = OpenFilePOptions.newBuilder().build();
      InStreamOptions options = new InStreamOptions(dummyStatus, readOptions, S_CONF, mContext);
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(new BlockInfo());
      when(mContext.getCachedWorkers()).thenReturn(
          Lists.newArrayList(new BlockWorkerInfo(remote, 100, 0),
              new BlockWorkerInfo(local, 100, 0)));
  
      when(mContext.getNodeLocalWorker()).thenReturn(local);
      when(mContext.hasProcessLocalWorker()).thenReturn(true);
      Optional<BlockWorker> blockWorker = Optional.of(Mockito.mock(BlockWorker.class));
      when(mContext.getProcessLocalWorker()).thenReturn(blockWorker);
  
      BlockInStream stream = mBlockStore.getInStream(BLOCK_ID, options);
      assertEquals(local, stream.getAddress());
      assertEquals(BlockWorkerDataReader.Factory.class.getName(),
          stream.getDataReaderFactory().getClass().getName());
    }
  
    @Test
    public void getInStreamInAlluxioOnlyFallbackToAvailableWorker() throws Exception {
      int workerCount = 4;
      boolean persisted = false;
      int[] blockLocations = new int[]{2, 3};
      Map<Integer, Long> failedWorkers = ImmutableMap.of(
          0, 3L,
          1, 1L,
          3, 2L);
      int expectedWorker = 2;
      testGetInStreamFallback(workerCount, persisted, blockLocations, failedWorkers, expectedWorker);
    }
  
    @Test
    public void getInStreamPersistedAndInAlluxioFallbackToUFS() throws Exception {
      int workerCount = 3;
      boolean persisted = true;
      int[] blockLocations = new int[]{0, 2};
      Map<Integer, Long> failedWorkers = ImmutableMap.of(
          0, 5L,
          2, 2L);
      int expectedWorker = 1;
      testGetInStreamFallback(workerCount, persisted, blockLocations, failedWorkers, expectedWorker);
    }
  
    @Test
    public void getInStreamPersistedFallbackToLeastRecentlyFailed() throws Exception {
      int workerCount = 3;
      boolean persisted = true;
      int[] blockLocations = new int[0];
      Map<Integer, Long> failedWorkers = ImmutableMap.of(
          0, 5L,
          1, 1L,
          2, 2L);
      int expectedWorker = 1;
      testGetInStreamFallback(workerCount, persisted, blockLocations, failedWorkers, expectedWorker);
    }
  
    @Test
    public void getInStreamInAlluxioOnlyFallbackToLeastRecentlyFailed() throws Exception {
      int workerCount = 5;
      boolean persisted = false;
      int[] blockLocations = new int[]{1, 2, 3};
      Map<Integer, Long> failedWorkers = ImmutableMap.of(
          0, 5L,
          1, 3L,
          2, 2L,
          3, 4L,
          4, 1L);
      int expectedWorker = 2;
      testGetInStreamFallback(workerCount, persisted, blockLocations, failedWorkers, expectedWorker);
    }
  
    @Test
    public void getInStreamInAlluxioWhenCreateStreamIsFailed() throws Exception {
      int workerCount = 5;
      boolean persisted = false;
      int[] blockLocations = new int[]{2, 3, 4};
      Map<Integer, Long> failedWorkers = ImmutableMap.of(
              0, 3L,
              1, 1L,
              3, 2L);
      int expectedWorker = 2;
      WorkerNetAddress[] workers = new WorkerNetAddress[workerCount];
      for (int i = 0; i < workers.length - 1; i++) {
        workers[i] = new WorkerNetAddress().setHost(String.format("worker-%d", i));
      }
      workers[workers.length - 1] = new WorkerNetAddress().setHost(WORKER_HOSTNAME_LOCAL);
      when(mContext.acquireBlockWorkerClient(WORKER_NET_ADDRESS_LOCAL))
          .thenThrow(new UnavailableException("failed to connect to "
              + WORKER_NET_ADDRESS_LOCAL.getHost()));
      BlockInfo info = new BlockInfo().setBlockId(BLOCK_ID)
          .setLocations(Arrays.stream(blockLocations).mapToObj(x ->
              new BlockLocation().setWorkerAddress(workers[x])).collect(Collectors.toList()));
      URIStatus dummyStatus =
          new URIStatus(new FileInfo().setPersisted(persisted)
              .setBlockIds(Collections.singletonList(BLOCK_ID))
              .setFileBlockInfos(Collections.singletonList(new FileBlockInfo().setBlockInfo(info))));
      BlockLocationPolicy mockPolicy = mock(BlockLocationPolicy.class);
      when(mockPolicy.getWorker(any())).thenAnswer(arg -> arg
          .getArgument(0, GetWorkerOptions.class).getBlockWorkerInfos().iterator().next()
          .getNetAddress());
      InStreamOptions options =
          new InStreamOptions(dummyStatus, FileSystemOptionsUtils.openFileDefaults(S_CONF),
              S_CONF, mContext);
      options.setUfsReadLocationPolicy(mockPolicy);
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(info);
      when(mContext.getCachedWorkers()).thenReturn(
          Arrays.stream(workers)
              .map(x -> new BlockWorkerInfo(x, -1, -1)).collect((Collectors.toList())));
      Map<WorkerNetAddress, Long> failedWorkerAddresses = failedWorkers.entrySet().stream()
          .map(x -> new AbstractMap.SimpleImmutableEntry<>(workers[x.getKey()], x.getValue()))
          .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
  
      BlockInStream inStream = null;
      int i = 2;
      while (i-- > 0) {
        try {
          inStream = mBlockStore.getInStream(BLOCK_ID, options,
                  failedWorkerAddresses);
        } catch (Exception e) {
          //do nothing
        }
      }
      Objects.requireNonNull(inStream);
      assertEquals(workers[expectedWorker], inStream.getAddress());
    }
  
    private void testGetInStreamFallback(int workerCount, boolean isPersisted, int[] blockLocations,
          Map<Integer, Long> failedWorkers, int expectedWorker) throws Exception {
      WorkerNetAddress[] workers = new WorkerNetAddress[workerCount];
      Arrays.setAll(workers, i -> new WorkerNetAddress().setHost(String.format("worker-%d", i)));
      BlockInfo info = new BlockInfo().setBlockId(BLOCK_ID)
          .setLocations(Arrays.stream(blockLocations).mapToObj(x ->
              new BlockLocation().setWorkerAddress(workers[x])).collect(Collectors.toList()));
      URIStatus dummyStatus =
          new URIStatus(new FileInfo().setPersisted(isPersisted)
              .setBlockIds(Collections.singletonList(BLOCK_ID))
              .setFileBlockInfos(Collections.singletonList(new FileBlockInfo().setBlockInfo(info))));
      BlockLocationPolicy mockPolicy = mock(BlockLocationPolicy.class);
      when(mockPolicy.getWorker(any())).thenAnswer(arg -> Optional.ofNullable(arg
          .getArgument(0, GetWorkerOptions.class).getBlockWorkerInfos().iterator().next()
          .getNetAddress()));
      InStreamOptions options =
          new InStreamOptions(dummyStatus, FileSystemOptionsUtils.openFileDefaults(S_CONF),
              S_CONF, mContext);
      options.setUfsReadLocationPolicy(mockPolicy);
      when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(info);
      when(mContext.getCachedWorkers()).thenReturn(Arrays.stream(workers)
          .map(x -> new BlockWorkerInfo(x, -1, -1)).collect((Collectors.toList())));
      Map<WorkerNetAddress, Long> failedWorkerAddresses = failedWorkers.entrySet().stream()
          .map(x -> new AbstractMap.SimpleImmutableEntry<>(workers[x.getKey()], x.getValue()))
          .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
  
      BlockInStream inStream = mBlockStore.getInStream(BLOCK_ID, options, failedWorkerAddresses);
  
      assertEquals(workers[expectedWorker], inStream.getAddress());
    }
  }
  