alluxio-master-2.x/core/client/fs/src/test/java/alluxio/client/block/stream/GrpcBlockingStreamTest.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.client.block.stream;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import alluxio.Constants;
import alluxio.exception.status.CancelledException;
import alluxio.exception.status.DeadlineExceededException;
import alluxio.exception.status.UnauthenticatedException;
import alluxio.grpc.WriteRequest;
import alluxio.grpc.WriteResponse;
import alluxio.util.ThreadFactoryUtils;

import io.grpc.Status;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.ClientResponseObserver;
import io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Test;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Stream;

/**
 * Tests for {@link GrpcBlockingStream}.
 */
public final class GrpcBlockingStreamTest {
  private static final int BUFFER_SIZE = 5;
  private static final long TIMEOUT = 10 * Constants.SECOND_MS;
  private static final long SHORT_TIMEOUT = Constants.SECOND / 2;
  private static final ExecutorService EXECUTOR = Executors.newFixedThreadPool(4,
      ThreadFactoryUtils.build("test-executor-%d", true));

  private static final String TEST_MESSAGE = "test message";

  private BlockWorkerClient mClient;
  private ClientCallStreamObserver<WriteRequest> mRequestObserver;
  private ClientResponseObserver<WriteRequest, WriteResponse> mResponseObserver;
  private GrpcBlockingStream<WriteRequest, WriteResponse> mStream;
  private Runnable mOnReadyHandler;

  /**
   * Set up gRPC interface mocks.
   */
  @Before
  public void before() {
    mClient = mock(BlockWorkerClient.class);
    mRequestObserver = mock(ClientCallStreamObserver.class);
    when(mClient.writeBlock(any(StreamObserver.class))).thenAnswer((args) -> {
      mResponseObserver = args.getArgument(0, ClientResponseObserver.class);
      return mRequestObserver;
    });
    when(mRequestObserver.isReady()).thenReturn(true);
    mStream = new GrpcBlockingStream<>(mClient::writeBlock, BUFFER_SIZE, TEST_MESSAGE);
  }

  /**
   * Checks send request is called on the request observer.
   */
  @Test
  public void send() throws Exception {
    WriteRequest request = WriteRequest.newBuilder().build();

    mStream.send(request, TIMEOUT);

    verify(mRequestObserver).onNext(request);
  }

  /**
   * Checks response posted on the response observer is received.
   */
  @Test
  public void receive() throws Exception {
    WriteResponse response = WriteResponse.newBuilder().build();
    mResponseObserver.onNext(response);

    WriteResponse actualResponse = mStream.receive(TIMEOUT);

    assertEquals(response, actualResponse);
  }

  /**
   * Checks onCompleted is called on request observer upon close.
   */
  @Test
  public void close() throws Exception {
    mStream.close();

    assertTrue(mStream.isClosed());
    assertFalse(mStream.isOpen());
    verify(mRequestObserver).onCompleted();
  }

  /**
   * Checks cancel is called on request observer upon cancel.
   */
  @Test
  public void cancel() throws Exception {
    mStream.cancel();

    assertTrue(mStream.isCanceled());
    assertFalse(mStream.isOpen());
    verify(mRequestObserver).cancel(any(String.class), eq(null));
  }

  /**
   * Checks onCompleted posted on the response observer is received.
   */
  @Test
  public void onCompleted() throws Exception {
    mResponseObserver.onCompleted();

    WriteResponse actualResponse = mStream.receive(TIMEOUT);
    assertNull(actualResponse);
  }

  /**
   * Checks expected error during send.
   */
  @Test
  public void sendError() throws Exception {
    mResponseObserver.onError(Status.UNAUTHENTICATED.asRuntimeException());
    Exception e = assertThrows(UnauthenticatedException.class,
        () -> mStream.send(WriteRequest.newBuilder().build(), TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks send fails after stream is closed.
   */
  @Test
  public void sendFailsAfterClosed() throws Exception {
    mStream.close();
    Exception e = assertThrows(CancelledException.class,
        () -> mStream.send(WriteRequest.newBuilder().build(), TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks send fails after stream is canceled.
   */
  @Test
  public void sendFailsAfterCanceled() throws Exception {
    mStream.cancel();
    Exception e = assertThrows(CancelledException.class, () ->
            mStream.send(WriteRequest.newBuilder().build(), TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks receive fails after stream is canceled.
   */
  @Test
  public void receiveFailsAfterCanceled() throws Exception {
    mStream.cancel();
    Exception e = assertThrows(CancelledException.class, () ->
        mStream.receive(TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks expected error during receive.
   */
  @Test
  public void receiveError() throws Exception {
    mResponseObserver.onError(Status.UNAUTHENTICATED.asRuntimeException());
    Exception e = assertThrows(UnauthenticatedException.class, () ->
        mStream.receive(TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks send fails after timeout waiting for stream to be ready.
   */
  @Test
  public void sendFailsAfterTimeout() throws Exception {
    when(mRequestObserver.isReady()).thenReturn(false);
    Exception e = assertThrows(DeadlineExceededException.class, () ->
        mStream.send(WriteRequest.newBuilder().build(), SHORT_TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks receive fails after timeout waiting for message from response stream.
   */
  @Test
  public void receiveFailsAfterTimeout() throws Exception {
    Exception e = assertThrows(DeadlineExceededException.class, () ->
        mStream.receive(SHORT_TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks send after stream is ready.
   */
  @Test
  public void sendAfterStreamReady() throws Exception {
    when(mRequestObserver.isReady()).thenReturn(false);
    doAnswer((args) -> {
      mOnReadyHandler = args.getArgument(0, Runnable.class);
      return null;
    }).when(mRequestObserver).setOnReadyHandler(any(Runnable.class));
    mResponseObserver.beforeStart(mRequestObserver);
    EXECUTOR.submit(() -> {
      try {
        // notify ready after a short period of time
        Thread.sleep(SHORT_TIMEOUT);
        when(mRequestObserver.isReady()).thenReturn(true);
        mOnReadyHandler.run();
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
    });
    WriteRequest request = WriteRequest.newBuilder().build();

    mStream.send(request, TIMEOUT);

    verify(mRequestObserver).onNext(request);
  }

  /**
   * Checks receive after response arrives.
   */
  @Test
  public void receiveAfterResponseArrives() throws Exception {
    WriteResponse response = WriteResponse.newBuilder().build();
    EXECUTOR.submit(() -> {
      try {
        // push response after a short period of time
        Thread.sleep(SHORT_TIMEOUT);
        mResponseObserver.onNext(response);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
    });

    WriteResponse actualResponse = mStream.receive(TIMEOUT);

    assertEquals(response, actualResponse);
  }

  /**
   * Checks receive responses more than buffer size in order.
   */
  @Test
  public void receiveMoreThanBufferSize() throws Exception {
    WriteResponse[] responses = Stream.generate(() -> WriteResponse.newBuilder().build())
        .limit(BUFFER_SIZE * 2).toArray(WriteResponse[]::new);

    EXECUTOR.submit(() -> {
      for (WriteResponse response : responses) {
        mResponseObserver.onNext(response);
      }
    });
    Thread.sleep(SHORT_TIMEOUT);

    for (WriteResponse response : responses) {
      WriteResponse actualResponse = mStream.receive(TIMEOUT);
      assertEquals(response, actualResponse);
    }
  }

  /**
   * Checks receive fails immediately upon error even if buffer is full.
   */
  @Test
  public void receiveErrorWhenBufferFull() throws Exception {
    WriteResponse[] responses = Stream.generate(() -> WriteResponse.newBuilder().build())
        .limit(BUFFER_SIZE).toArray(WriteResponse[]::new);
    for (WriteResponse response : responses) {
      mResponseObserver.onNext(response);
    }
    mResponseObserver.onError(Status.UNAUTHENTICATED.asRuntimeException());

    Exception e = assertThrows(UnauthenticatedException.class, () -> {
      for (WriteResponse response : responses) {
        WriteResponse actualResponse = mStream.receive(TIMEOUT);
        assertEquals(response, actualResponse);
      }
    });
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }

  /**
   * Checks waitForComplete succeed after onCompleted is triggered on response stream.
   */
  @Test
  public void waitForComplete() throws Exception {
    WriteResponse[] responses = Stream.generate(() -> WriteResponse.newBuilder().build())
        .limit(BUFFER_SIZE * 2).toArray(WriteResponse[]::new);

    EXECUTOR.submit(() -> {
      for (WriteResponse response : responses) {
        mResponseObserver.onNext(response);
      }
      try {
        Thread.sleep(SHORT_TIMEOUT);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
      mResponseObserver.onCompleted();
    });

    WriteResponse actualResponse = mStream.receive(TIMEOUT);
    assertEquals(responses[0], actualResponse);

    mStream.waitForComplete(TIMEOUT);

    actualResponse = mStream.receive(TIMEOUT);
    assertEquals(null, actualResponse);
  }

  /**
   * Checks waitForComplete fails after times out.
   */
  @Test
  public void waitForCompleteTimeout() throws Exception {
    WriteResponse[] responses = Stream.generate(() -> WriteResponse.newBuilder().build())
        .limit(BUFFER_SIZE * 2).toArray(WriteResponse[]::new);

    EXECUTOR.submit(() -> {
      for (WriteResponse response : responses) {
        mResponseObserver.onNext(response);
      }
      try {
        Thread.sleep(TIMEOUT);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
      mResponseObserver.onCompleted();
    });

    WriteResponse actualResponse = mStream.receive(SHORT_TIMEOUT);
    assertEquals(responses[0], actualResponse);

    Exception e = assertThrows(DeadlineExceededException.class, () ->
        mStream.waitForComplete(SHORT_TIMEOUT));
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }
}
