StreamBodyConsumerTestBase.java

/*
 * Copyright © 2015 Cask Data, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package co.cask.cdap.data.stream.service.upload;

import co.cask.cdap.proto.Id;
import co.cask.common.io.ByteBufferInputStream;
import co.cask.http.AbstractHttpResponder;
import co.cask.http.BodyConsumer;
import co.cask.http.BodyProducer;
import co.cask.http.ChunkResponder;
import co.cask.http.HttpResponder;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.io.Files;
import com.google.common.io.InputSupplier;
import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBuffers;
import org.jboss.netty.handler.codec.http.HttpResponseStatus;
import org.junit.Assert;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Unit test base for various stream {@link BodyConsumer} that supports the stream batch endpoint.
 */
public abstract class StreamBodyConsumerTestBase {

  @ClassRule
  public static final TemporaryFolder TMP_FOLDER = new TemporaryFolder();

  protected abstract ContentInfo generateFile(int recordCount) throws IOException;

  protected abstract BodyConsumer createBodyConsumer(ContentWriterFactory contentWriterFactory);

  @Test
  public void testChunkedContent() throws Exception {
    int recordCount = 1000;
    // Generate a file for upload
    ContentInfo contentInfo = generateFile(recordCount);

    final Map<String, String> contentHeaders = Maps.newHashMap();
    final TestContentWriter contentWriter = new TestContentWriter();
    BodyConsumer bodyConsumer = createBodyConsumer(new ContentWriterFactory() {
      @Override
      public Id.Stream getStream() {
        return Id.Stream.from("test_namespace", "test-stream");
      }

      @Override
      public ContentWriter create(Map<String, String> headers) throws IOException {
        contentHeaders.putAll(headers);
        return contentWriter;
      }
    });
    TestHttpResponder responder = new TestHttpResponder();

    // Feed the file content in small chunk
    sendChunks(contentInfo.getContentSupplier(), 10, bodyConsumer, responder);

    // Verify the processing is completed correctly
    Assert.assertTrue(contentWriter.waitForClose(5, TimeUnit.SECONDS));
    Assert.assertEquals(HttpResponseStatus.OK, responder.getResponseStatus());
    Assert.assertEquals(recordCount, contentWriter.getEvents());
    Assert.assertTrue(contentInfo.verify(contentHeaders, new InputSupplier<InputStream>() {
      @Override
      public InputStream getInput() throws IOException {
        return new ByteBufferInputStream(contentWriter.getContent().duplicate());
      }
    }));
  }

  /**
   * Sends content as provided by the given {@link InputSupplier} as small chunks to the given
   * {@link BodyConsumer}.
   */
  private void sendChunks(InputSupplier<? extends InputStream> inputSupplier, int chunkSize,
                          BodyConsumer bodyConsumer, HttpResponder responder) throws IOException {
    InputStream input = inputSupplier.getInput();
    try {
      byte[] bytes = new byte[chunkSize];
      int len = input.read(bytes);
      while (len >= 0) {
        bodyConsumer.chunk(ChannelBuffers.copiedBuffer(bytes, 0, len), responder);
        len = input.read(bytes);
      }
      bodyConsumer.finished(responder);
    } catch (Exception e) {
      bodyConsumer.handleError(e);
    } finally {
      input.close();
    }
  }

  /**
   * An interface for sub-classes to implement to provide information about content being provided to body consumer.
   */
  protected interface ContentInfo {
    /**
     * Returns an {@link InputSupplier} that provides the content.
     */
    InputSupplier<? extends InputStream> getContentSupplier();

    /**
     * Verify the uploaded content after being handled by the body consumer
     * created by the {@link #createBodyConsumer(ContentWriterFactory)} method.
     */
    boolean verify(Map<String, String> headers,
                   InputSupplier<? extends InputStream> contentSupplier) throws IOException;
  }

  /**
   * An implementation of {@link ContentInfo} that reads content from the given file.
   */
  protected abstract class FileContentInfo implements ContentInfo {

    private final File file;

    protected FileContentInfo(File file) {
      this.file = file;
    }

    @Override
    public InputSupplier<? extends InputStream> getContentSupplier() {
      return Files.newInputStreamSupplier(file);
    }
  }

  /**
   * A {@link ContentWriter} for testing. It keeps all content written in memory.
   */
  private static class TestContentWriter implements ContentWriter {
    private final List<ChannelBuffer> contents;
    private final CountDownLatch completion;
    private int events;

    public TestContentWriter() {
      this.contents = Lists.newLinkedList();
      this.completion = new CountDownLatch(1);
    }

    @Override
    public void append(ByteBuffer body, boolean immutable) throws IOException {
      contents.add(immutable ? ChannelBuffers.wrappedBuffer(body) : ChannelBuffers.copiedBuffer(body));
      events++;
    }

    @Override
    public void appendAll(Iterator<ByteBuffer> bodies, boolean immutable) throws IOException {
      while (bodies.hasNext()) {
        append(bodies.next(), immutable);
      }
    }

    @Override
    public void close() throws IOException {
      completion.countDown();
    }

    public boolean waitForClose(long timeout, TimeUnit unit) throws InterruptedException {
      return completion.await(timeout, unit);
    }

    public ByteBuffer getContent() {
      return ChannelBuffers.wrappedBuffer(contents.toArray(new ChannelBuffer[contents.size()])).toByteBuffer();
    }

    @Override
    public void cancel() {
      contents.clear();
      completion.countDown();
    }

    public int getEvents() {
      return events;
    }
  }

  /**
   * A {@link HttpResponder} for testing. It only saved the first response status event sent.
   */
  private static class TestHttpResponder extends AbstractHttpResponder {
    private final AtomicReference<HttpResponseStatus> responseStatus = new AtomicReference<>();

    @Override
    public ChunkResponder sendChunkStart(HttpResponseStatus status, Multimap<String, String> headers) {
      // Not used in test
      return null;
    }

    @Override
    public void sendContent(HttpResponseStatus status, ChannelBuffer content,
                            String contentType, Multimap<String, String> headers) {
      responseStatus.compareAndSet(null, status);
    }

    @Override
    public void sendFile(File file, Multimap<String, String> headers) {
      // Not used in test
    }

    @Override
    public void sendContent(HttpResponseStatus httpResponseStatus,
                            BodyProducer bodyProducer, Multimap<String, String> multimap) {
      // Not used in test
    }

    public HttpResponseStatus getResponseStatus() {
      return responseStatus.get();
    }
  }
}