core/src/test/java/com/google/net/stubby/newtransport/GrpcDeframerTest.java

package com.google.net.stubby.newtransport;

import static com.google.net.stubby.newtransport.TransportFrameUtil.CONTEXT_VALUE_FRAME;
import static com.google.net.stubby.newtransport.TransportFrameUtil.PAYLOAD_FRAME;
import static com.google.net.stubby.newtransport.TransportFrameUtil.STATUS_FRAME;
import static java.nio.charset.StandardCharsets.UTF_8;
import static org.junit.Assert.assertEquals;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyInt;
import static org.mockito.Matchers.anyString;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.google.common.io.ByteStreams;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.SettableFuture;
import com.google.net.stubby.Status;
import com.google.net.stubby.transport.Transport;
import com.google.protobuf.ByteString;

import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;

import javax.annotation.Nullable;

/**
 * Tests for {@link GrpcDeframer}.
 */
@RunWith(JUnit4.class)
public class GrpcDeframerTest {
  private static final String KEY = "key";
  private static final String MESSAGE = "hello world";
  private static final ByteString MESSAGE_BSTR = ByteString.copyFromUtf8(MESSAGE);
  private static final Transport.Code STATUS_CODE = Transport.Code.CANCELLED;

  private GrpcDeframer reader;

  private Transport.ContextValue contextProto;

  private StubDecompressor decompressor;

  @Mock
  private StreamListener listener;

  private SettableFuture<Void> contextFuture;

  private SettableFuture<Void> messageFuture;

  @Before
  public void setup() {
    MockitoAnnotations.initMocks(this);

    contextFuture = SettableFuture.create();
    messageFuture = SettableFuture.create();
    when(listener.contextRead(anyString(), any(InputStream.class), anyInt())).thenReturn(
        contextFuture);
    when(listener.messageRead(any(InputStream.class), anyInt())).thenReturn(messageFuture);

    decompressor = new StubDecompressor();
    reader = new GrpcDeframer(decompressor, listener, MoreExecutors.directExecutor());

    contextProto = Transport.ContextValue.newBuilder().setKey(KEY).setValue(MESSAGE_BSTR).build();
  }

  @Test
  public void contextShouldCallTarget() throws Exception {
    decompressor.init(contextFrame());
    reader.deframe(Buffers.empty(), false);
    verifyContext();
    verifyNoPayload();
    verifyNoStatus();
  }

  @Test
  public void contextWithEndOfStreamShouldNotifyStatus() throws Exception {
    decompressor.init(contextFrame());
    reader.deframe(Buffers.empty(), true);
    verifyContext();
    verifyNoPayload();
    verifyNoStatus();

    contextFuture.set(null);
    verifyStatus(Transport.Code.OK);
  }

  @Test
  public void payloadShouldCallTarget() throws Exception {
    decompressor.init(payloadFrame());
    reader.deframe(Buffers.empty(), false);
    verifyNoContext();
    verifyPayload();
    verifyNoStatus();
  }

  @Test
  public void payloadWithEndOfStreamShouldNotifyStatus() throws Exception {
    decompressor.init(payloadFrame());
    reader.deframe(Buffers.empty(), true);
    verifyNoContext();
    verifyPayload();
    verifyNoStatus();

    messageFuture.set(null);
    verifyStatus(Transport.Code.OK);
  }

  @Test
  public void statusShouldCallTarget() throws Exception {
    decompressor.init(statusFrame());
    reader.deframe(Buffers.empty(), false);
    verifyNoContext();
    verifyNoPayload();
    verifyStatus();
  }

  @Test
  public void statusWithEndOfStreamShouldNotifyStatusOnce() throws Exception {
    decompressor.init(statusFrame());
    reader.deframe(Buffers.empty(), true);
    verifyNoContext();
    verifyNoPayload();
    verifyStatus();
  }

  @Test
  public void multipleFramesShouldCallTarget() throws Exception {
    ByteArrayOutputStream os = new ByteArrayOutputStream();
    DataOutputStream dos = new DataOutputStream(os);

    // Write a context frame.
    writeFrame(CONTEXT_VALUE_FRAME, contextProto.toByteArray(), dos);

    // Write a payload frame.
    writeFrame(PAYLOAD_FRAME, MESSAGE_BSTR.toByteArray(), dos);

    // Write a status frame.
    byte[] statusBytes = new byte[] {0, (byte) STATUS_CODE.getNumber()};
    writeFrame(STATUS_FRAME, statusBytes, dos);

    // Now write the complete frame: compression header followed by the 3 message frames.
    dos.close();
    byte[] bodyBytes = os.toByteArray();

    decompressor.init(bodyBytes);
    reader.deframe(Buffers.empty(), false);

    // Verify that all callbacks were called.
    verifyContext();
    verifyNoPayload();
    verifyNoStatus();

    contextFuture.set(null);
    verifyPayload();
    verifyNoStatus();

    messageFuture.set(null);
    verifyStatus();
  }

  @Test
  public void partialFrameShouldSucceed() throws Exception {
    byte[] frame = payloadFrame();

    // Create a buffer that contains 2 payload frames.
    byte[] fullBuffer = Arrays.copyOf(frame, frame.length * 2);
    System.arraycopy(frame, 0, fullBuffer, frame.length, frame.length);

    // Use only a portion of the frame. Should not call the listener.
    int startIx = 0;
    int endIx = 10;
    byte[] chunk = Arrays.copyOfRange(fullBuffer, startIx, endIx);
    decompressor.init(chunk);
    reader.deframe(Buffers.empty(), false);
    verifyNoContext();
    verifyNoPayload();
    verifyNoStatus();

    // Supply the rest of the frame and a portion of a second frame. Should call the listener.
    startIx = endIx;
    endIx = startIx + frame.length;
    chunk = Arrays.copyOfRange(fullBuffer, startIx, endIx);
    decompressor.init(chunk);
    reader.deframe(Buffers.empty(), false);
    verifyNoContext();
    verifyPayload();
    verifyNoStatus();
  }

  private void verifyContext() {
    ArgumentCaptor<InputStream> captor = ArgumentCaptor.forClass(InputStream.class);
    verify(listener).contextRead(eq(KEY), captor.capture(), eq(MESSAGE.length()));
    assertEquals(MESSAGE, readString(captor.getValue(), MESSAGE.length()));
  }

  private void verifyPayload() {
    ArgumentCaptor<InputStream> captor = ArgumentCaptor.forClass(InputStream.class);
    verify(listener).messageRead(captor.capture(), eq(MESSAGE.length()));
    assertEquals(MESSAGE, readString(captor.getValue(), MESSAGE.length()));
  }

  private String readString(InputStream in, int length) {
    try {
      byte[] bytes = new byte[length];
      ByteStreams.readFully(in, bytes);
      return new String(bytes, UTF_8);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  private void verifyStatus() {
    verifyStatus(Transport.Code.CANCELLED);
  }

  private void verifyStatus(Transport.Code code) {
    ArgumentCaptor<Status> captor = ArgumentCaptor.forClass(Status.class);
    verify(listener).closed(captor.capture());
    assertEquals(code, captor.getValue().getCode());
  }

  private void verifyNoContext() {
    verify(listener, never()).contextRead(any(String.class), any(InputStream.class), anyInt());
  }

  private void verifyNoPayload() {
    verify(listener, never()).messageRead(any(InputStream.class), anyInt());
  }

  private void verifyNoStatus() {
    verify(listener, never()).closed(any(Status.class));
  }

  private byte[] contextFrame() throws IOException {
    return frame(CONTEXT_VALUE_FRAME, contextProto.toByteArray());
  }

  private static byte[] payloadFrame() throws IOException {
    return frame(PAYLOAD_FRAME, MESSAGE_BSTR.toByteArray());
  }

  private static byte[] statusFrame() throws IOException {
    byte[] bytes = new byte[] {0, (byte) STATUS_CODE.getNumber()};
    return frame(STATUS_FRAME, bytes);
  }

  private static byte[] frame(int frameType, byte[] data) throws IOException {
    ByteArrayOutputStream bos = new ByteArrayOutputStream();
    OutputStream os = bos;
    DataOutputStream dos = new DataOutputStream(os);
    writeFrame(frameType, data, dos);
    dos.close();
    return bos.toByteArray();
  }

  private static void writeFrame(int frameType, byte[] data, DataOutputStream out)
      throws IOException {
    out.write(frameType);
    out.writeInt(data.length);
    out.write(data);
  }

  private static final class StubDecompressor implements Decompressor {
    byte[] bytes;
    int offset;

    void init(byte[] bytes) {
      this.bytes = bytes;
      this.offset = 0;
    }

    @Override
    public void decompress(Buffer data) {
      // Do nothing.
    }

    @Override
    public void close() {
      // Do nothing.
    }

    @Override
    @Nullable
    public Buffer readBytes(int length) {
      length = Math.min(length, bytes.length - offset);
      if (length == 0) {
        return null;
      }

      Buffer buffer = Buffers.wrap(ByteString.copyFrom(bytes, offset, length));
      offset += length;
      return buffer;
    }
  }
}