ChunkedStreamTest.java

package org.jooby.internal.undertow;

import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.isA;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;

import org.jooby.test.MockUnit;
import org.jooby.test.MockUnit.Block;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;

import io.undertow.connector.ByteBufferPool;
import io.undertow.connector.PooledByteBuffer;
import io.undertow.io.IoCallback;
import io.undertow.io.Sender;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.ServerConnection;
import io.undertow.util.HeaderMap;
import io.undertow.util.Headers;

@RunWith(PowerMockRunner.class)
@PrepareForTest({ChunkedStream.class, HttpServerExchange.class, HeaderMap.class })
public class ChunkedStreamTest {

  int bufferSize = 10;
  byte[] bytes = "bytes".getBytes();
  ByteBuffer buffer = ByteBuffer.wrap(bytes);

  private Block exchange = unit -> {
    ByteBufferPool bpool = unit.mock(ByteBufferPool.class);
    expect(bpool.allocate()).andReturn(unit.get(PooledByteBuffer.class));

    ServerConnection conn = unit.mock(ServerConnection.class);
    expect(conn.getByteBufferPool()).andReturn(bpool);
    expect(conn.getBufferSize()).andReturn(bufferSize);

    HttpServerExchange exchange = unit.get(HttpServerExchange.class);
    expect(exchange.getResponseSender()).andReturn(unit.get(Sender.class));
    expect(exchange.getConnection()).andReturn(conn);
  };

  private Block noIoThread = unit -> {
    HttpServerExchange exchange = unit.get(HttpServerExchange.class);
    expect(exchange.isInIoThread()).andReturn(false);
  };

  private Block ioThread = unit -> {
    HttpServerExchange exchange = unit.get(HttpServerExchange.class);
    expect(exchange.isInIoThread()).andReturn(true);
    expect(exchange.dispatch(isA(ChunkedStream.class))).andReturn(exchange);
  };

  private Block nolen = unit -> {
    HeaderMap headers = unit.mock(HeaderMap.class);
    expect(headers.contains(Headers.CONTENT_LENGTH)).andReturn(false);
    expect(headers.put(Headers.CONTENT_LENGTH, bytes.length)).andReturn(headers);
    expect(headers.remove(Headers.TRANSFER_ENCODING)).andReturn(null);

    HttpServerExchange exchange = unit.get(HttpServerExchange.class);
    expect(exchange.getResponseHeaders()).andReturn(headers);
  };

  private Block len = unit -> {
    HeaderMap headers = unit.mock(HeaderMap.class);
    expect(headers.contains(Headers.CONTENT_LENGTH)).andReturn(true);

    HttpServerExchange exchange = unit.get(HttpServerExchange.class);
    expect(exchange.getResponseHeaders()).andReturn(headers);
  };

  private Block transferEncoding = unit -> {
    HeaderMap headers = unit.mock(HeaderMap.class);
    expect(headers.contains(Headers.CONTENT_LENGTH)).andReturn(false);
    expect(headers.put(Headers.TRANSFER_ENCODING, "chunked")).andReturn(headers);

    HttpServerExchange exchange = unit.get(HttpServerExchange.class);
    expect(exchange.getResponseHeaders()).andReturn(headers);
  };

  private Block sendChunk = unit -> {
    Sender sender = unit.get(Sender.class);

    sender.send(eq(buffer), isA(ChunkedStream.class));
  };

  private Block readChunk = unit -> {
    ReadableByteChannel channel = unit.get(ReadableByteChannel.class);
    expect(channel.read(buffer)).andReturn(bytes.length);
  };

  private Block readNoChunk = unit -> {
    ReadableByteChannel channel = unit.get(ReadableByteChannel.class);
    expect(channel.read(buffer)).andReturn(-1);
  };

  private Block readErrChunk = unit -> {
    ReadableByteChannel channel = unit.get(ReadableByteChannel.class);
    expect(channel.read(buffer)).andThrow(new IOException("intentional err"));
  };

  private Block readLargeChunk = unit -> {
    ReadableByteChannel channel = unit.get(ReadableByteChannel.class);
    expect(channel.read(buffer)).andReturn(bufferSize);
  };

  private Block pooled = unit -> {
    PooledByteBuffer pooled = unit.get(PooledByteBuffer.class);
    expect(pooled.getBuffer()).andReturn(buffer);
  };

  @Test
  public void sendNoIoThread() throws Exception {
    new MockUnit(ReadableByteChannel.class, HttpServerExchange.class, IoCallback.class,
        Sender.class, PooledByteBuffer.class)
        .expect(exchange)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readChunk)
        .expect(nolen)
        .expect(sendChunk)
        .run(unit -> {
          new ChunkedStream().send(unit.get(ReadableByteChannel.class),
              unit.get(HttpServerExchange.class), unit.get(IoCallback.class));
        });
  }

  @Test
  public void sendNoIoThreadWithLen() throws Exception {
    new MockUnit(ReadableByteChannel.class, HttpServerExchange.class, IoCallback.class,
        Sender.class, PooledByteBuffer.class)
        .expect(exchange)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readChunk)
        .expect(unit -> {
          HeaderMap headers = unit.mock(HeaderMap.class);
          expect(headers.contains(Headers.CONTENT_LENGTH)).andReturn(true);

          HttpServerExchange exchange = unit.get(HttpServerExchange.class);
          expect(exchange.getResponseHeaders()).andReturn(headers);
        })
        .expect(sendChunk)
        .run(unit -> {
          new ChunkedStream().send(unit.get(ReadableByteChannel.class),
              unit.get(HttpServerExchange.class), unit.get(IoCallback.class));
        });
  }

  @Test
  public void sendNoIoThread2Chunks() throws Exception {
    new MockUnit(ReadableByteChannel.class, HttpServerExchange.class, IoCallback.class,
        Sender.class, PooledByteBuffer.class)
        .expect(exchange)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readLargeChunk)
        .expect(transferEncoding)
        .expect(sendChunk)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readChunk)
        .expect(sendChunk)
        .run(unit -> {
          ChunkedStream chunkedStream = new ChunkedStream();
          chunkedStream.send(unit.get(ReadableByteChannel.class),
              unit.get(HttpServerExchange.class), unit.get(IoCallback.class));
          chunkedStream.onComplete(unit.get(HttpServerExchange.class), unit.get(Sender.class));
        });
  }

  @Test
  public void sendNoIoThread2ChunksWithLen() throws Exception {
    new MockUnit(ReadableByteChannel.class, HttpServerExchange.class, IoCallback.class,
        Sender.class, PooledByteBuffer.class)
        .expect(exchange)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readLargeChunk)
        .expect(len)
        .expect(sendChunk)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readChunk)
        .expect(sendChunk)
        .run(unit -> {
          ChunkedStream chunkedStream = new ChunkedStream();
          chunkedStream.send(unit.get(ReadableByteChannel.class),
              unit.get(HttpServerExchange.class), unit.get(IoCallback.class));
          chunkedStream.onComplete(unit.get(HttpServerExchange.class), unit.get(Sender.class));
        });
  }

  @Test
  public void sendIoThread() throws Exception {
    new MockUnit(ReadableByteChannel.class, HttpServerExchange.class, IoCallback.class,
        Sender.class, PooledByteBuffer.class)
        .expect(exchange)
        .expect(ioThread)
        .expect(pooled)
        .expect(readChunk)
        .expect(nolen)
        .expect(sendChunk)
        .run(unit -> {
          ChunkedStream chunkedStream = new ChunkedStream();
          chunkedStream.send(unit.get(ReadableByteChannel.class),
              unit.get(HttpServerExchange.class), unit.get(IoCallback.class));
          chunkedStream.run();
        });
  }

  @Test
  public void sendNoIoThreadWithErr() throws Exception {
    new MockUnit(ReadableByteChannel.class, HttpServerExchange.class, IoCallback.class,
        Sender.class, PooledByteBuffer.class)
        .expect(exchange)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readErrChunk)
        .expect(unit -> {
          PooledByteBuffer pooled = unit.get(PooledByteBuffer.class);
          pooled.close();
          IoCallback err = unit.get(IoCallback.class);
          err.onException(eq(unit.get(HttpServerExchange.class)), eq(unit.get(Sender.class)),
              isA(IOException.class));
        })
        .run(unit -> {
          new ChunkedStream().send(unit.get(ReadableByteChannel.class),
              unit.get(HttpServerExchange.class), unit.get(IoCallback.class));
        });
  }

  @Test
  public void sendDone() throws Exception {
    new MockUnit(ReadableByteChannel.class, HttpServerExchange.class, IoCallback.class,
        Sender.class, PooledByteBuffer.class)
        .expect(exchange)
        .expect(noIoThread)
        .expect(pooled)
        .expect(readNoChunk)
        .expect(unit -> {
          PooledByteBuffer pooled = unit.get(PooledByteBuffer.class);
          pooled.close();
          IoCallback success = unit.get(IoCallback.class);
          success.onComplete(eq(unit.get(HttpServerExchange.class)), eq(unit.get(Sender.class)));
        })
        .run(unit -> {
          new ChunkedStream().send(unit.get(ReadableByteChannel.class),
              unit.get(HttpServerExchange.class), unit.get(IoCallback.class));
        });
  }

}