TestMIMEIntegrationReaderWriter.java

/*
   Copyright (c) 2015 LinkedIn Corp.

   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 com.linkedin.multipart.integ;


import com.linkedin.common.callback.Callback;
import com.linkedin.data.ByteString;
import com.linkedin.multipart.MultiPartMIMEDataSourceWriter;
import com.linkedin.multipart.MultiPartMIMEInputStream;
import com.linkedin.multipart.MultiPartMIMEReader;
import com.linkedin.multipart.MultiPartMIMEReaderCallback;
import com.linkedin.multipart.MultiPartMIMEStreamRequestFactory;
import com.linkedin.multipart.MultiPartMIMEWriter;
import com.linkedin.multipart.SinglePartMIMEReaderCallback;
import com.linkedin.multipart.exceptions.MultiPartIllegalFormatException;
import com.linkedin.multipart.utils.MIMEDataPart;
import com.linkedin.r2.filter.R2Constants;
import com.linkedin.r2.message.Messages;
import com.linkedin.r2.message.RequestContext;
import com.linkedin.r2.message.rest.RestException;
import com.linkedin.r2.message.rest.RestResponse;
import com.linkedin.r2.message.rest.RestStatus;
import com.linkedin.r2.message.stream.StreamRequest;
import com.linkedin.r2.message.stream.StreamResponse;
import com.linkedin.r2.sample.Bootstrap;
import com.linkedin.r2.transport.common.StreamRequestHandler;
import com.linkedin.r2.transport.common.bridge.server.TransportDispatcher;
import com.linkedin.r2.transport.common.bridge.server.TransportDispatcherBuilder;
import com.linkedin.r2.transport.http.client.HttpClientFactory;

import com.google.common.collect.ImmutableList;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.URI;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.testng.Assert;
import org.testng.annotations.AfterSuite;
import org.testng.annotations.BeforeSuite;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;


/**
 * Represents integration tests that work between the {@link com.linkedin.multipart.MultiPartMIMEWriter} and
 * the {@link com.linkedin.multipart.MultiPartMIMEReader}
 *
 * @author Karim Vidhani
 */
public class TestMIMEIntegrationReaderWriter extends AbstractMIMEIntegrationStreamTest
{
  private static ScheduledExecutorService scheduledExecutorService;

  private static final URI SERVER_URI = URI.create("/pegasusMimeServer");
  private MIMEServerRequestHandler _mimeServerRequestHandler;

  byte[] _normalBodyData;
  Map<String, String> _normalBodyHeaders;

  byte[] _headerLessBodyData;

  Map<String, String> _bodyLessHeaders;

  MIMEDataPart _normalBody;
  MIMEDataPart _headerLessBody;
  MIMEDataPart _bodyLessBody;
  MIMEDataPart _purelyEmptyBody;

  @BeforeSuite
  public void dataSourceSetup()
  {
    scheduledExecutorService = Executors.newScheduledThreadPool(10);

    _normalBodyData = "some normal body that is relatively small".getBytes();
    _normalBodyHeaders = new HashMap<String, String>();
    _normalBodyHeaders.put("simpleheader", "simplevalue");

    //Second body has no headers
    _headerLessBodyData = "a body without headers".getBytes();

    //Third body has only headers
    _bodyLessHeaders = new HashMap<String, String>();
    _normalBodyHeaders.put("header1", "value1");
    _normalBodyHeaders.put("header2", "value2");
    _normalBodyHeaders.put("header3", "value3");

    _normalBody = new MIMEDataPart(ByteString.copy(_normalBodyData), _normalBodyHeaders);

    _headerLessBody = new MIMEDataPart(ByteString.copy(_headerLessBodyData), Collections.<String, String>emptyMap());

    _bodyLessBody = new MIMEDataPart(ByteString.empty(), _bodyLessHeaders);

    _purelyEmptyBody = new MIMEDataPart(ByteString.empty(), Collections.<String, String>emptyMap());
  }

  @AfterSuite
  public void shutDown()
  {
    scheduledExecutorService.shutdownNow();
  }

  @Override
  protected TransportDispatcher getTransportDispatcher()
  {
    _mimeServerRequestHandler = new MIMEServerRequestHandler();
    return new TransportDispatcherBuilder().addStreamHandler(SERVER_URI, _mimeServerRequestHandler).build();
  }

  @Override
  protected Map<String, String> getClientProperties()
  {
    Map<String, String> clientProperties = new HashMap<String, String>();
    clientProperties.put(HttpClientFactory.HTTP_REQUEST_TIMEOUT, "9000000");
    return clientProperties;
  }

  ///////////////////////////////////////////////////////////////////////////////////////

  @DataProvider(name = "eachSingleBodyDataSource")
  public Object[][] eachSingleBodyDataSource() throws Exception
  {
    return new Object[][]
        {
            {1, _normalBody},
            {10, _normalBody},
            {R2Constants.DEFAULT_DATA_CHUNK_SIZE, _normalBody},
            {1, _headerLessBody},
            {10, _headerLessBody},
            {R2Constants.DEFAULT_DATA_CHUNK_SIZE, _headerLessBody},
            {1, _bodyLessBody},
            {10, _bodyLessBody},
            {R2Constants.DEFAULT_DATA_CHUNK_SIZE, _bodyLessBody},
            {1, _purelyEmptyBody},
            {10, _purelyEmptyBody},
            {R2Constants.DEFAULT_DATA_CHUNK_SIZE, _purelyEmptyBody}
        };
  }

  @Test(dataProvider = "eachSingleBodyDataSource", enabled = false)
  public void testEachSingleBodyDataSource(final int chunkSize, final MIMEDataPart bodyPart) throws Exception
  {
    final MultiPartMIMEInputStream inputStreamDataSource =
        new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(bodyPart.getPartData().copyBytes()),
                                             scheduledExecutorService, bodyPart.getPartHeaders()).withWriteChunkSize(chunkSize).build();

    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder("some preamble", "").appendDataSource(inputStreamDataSource).build();

    executeRequestAndAssert(writer, ImmutableList.of(bodyPart));
  }

  @Test(dataProvider = "eachSingleBodyDataSource", enabled = false)
  public void testEachSingleBodyDataSourceMultipleTimes(final int chunkSize, final MIMEDataPart bodyPart)
      throws Exception
  {
    final List<MultiPartMIMEDataSourceWriter> dataSources = new ArrayList<MultiPartMIMEDataSourceWriter>();
    for (int i = 0; i < 4; i++)
    {
      final MultiPartMIMEInputStream inputStreamDataSource =
          new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(bodyPart.getPartData().copyBytes()),
                                               scheduledExecutorService, bodyPart.getPartHeaders()).withWriteChunkSize(chunkSize).build();
      dataSources.add(inputStreamDataSource);
    }

    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder("some preamble", "").appendDataSources(dataSources).build();

    executeRequestAndAssert(writer, ImmutableList.of(bodyPart, bodyPart, bodyPart, bodyPart));
  }

  ///////////////////////////////////////////////////////////////////////////////////////

  @DataProvider(name = "chunkSizes")
  public Object[][] chunkSizes() throws Exception
  {
    return new Object[][]
        {
            {1}, {R2Constants.DEFAULT_DATA_CHUNK_SIZE}
        };
  }

  @Test(dataProvider = "chunkSizes")
  public void testMultipleBodies(final int chunkSize) throws Exception
  {
    final MultiPartMIMEInputStream normalBodyInputStream =
        new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_normalBody.getPartData().copyBytes()),
                                             scheduledExecutorService, _normalBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();

    final MultiPartMIMEInputStream headerLessBodyInputStream =
        new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_headerLessBody.getPartData().copyBytes()),
                                             scheduledExecutorService, _headerLessBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();

    //Copying over empty ByteString, but let's keep things consistent.
    final MultiPartMIMEInputStream bodyLessBodyInputStream =
        new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_bodyLessBody.getPartData().copyBytes()),
                                             scheduledExecutorService, _bodyLessBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();

    final MultiPartMIMEInputStream purelyEmptyBodyInputStream =
        new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_purelyEmptyBody.getPartData().copyBytes()),
                                             scheduledExecutorService, _purelyEmptyBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();

    final MultiPartMIMEWriter writer =
        new MultiPartMIMEWriter.Builder("some preamble", "").appendDataSource(normalBodyInputStream)
            .appendDataSource(headerLessBodyInputStream).appendDataSource(bodyLessBodyInputStream)
            .appendDataSource(purelyEmptyBodyInputStream).build();

    executeRequestAndAssert(writer, ImmutableList.of(_normalBody, _headerLessBody, _bodyLessBody, _purelyEmptyBody));
  }

  @Test
  public void testEmptyEnvelope() throws Exception
  {
    //Javax mail does not support this, hence we can only test this using our writer
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder("some preamble", "").build();

    executeRequestAndAssert(writer, Collections.<MIMEDataPart>emptyList());
  }

  ///////////////////////////////////////////////////////////////////////////////////////

  //This method will execute a StreamRequest using a payload comprised of the provided MultiPartMIMEWriter. The request
  //makes its way to a server which will read all the parts one by one and store them in a list. Subsequently,
  //an assertion is performed to make sure that each part was read and saved properly on the server.
  private void executeRequestAndAssert(final MultiPartMIMEWriter requestWriter, final List<MIMEDataPart> expectedParts)
      throws Exception
  {
    final StreamRequest streamRequest =
        MultiPartMIMEStreamRequestFactory
            .generateMultiPartMIMEStreamRequest(Bootstrap.createHttpURI(PORT, SERVER_URI), "mixed", requestWriter,
                                                Collections.<String, String>emptyMap());

    final AtomicInteger status = new AtomicInteger(-1);
    final CountDownLatch latch = new CountDownLatch(1);
    Callback<StreamResponse> callback = expectSuccessCallback(latch, status, new HashMap<String, String>());
    _client.streamRequest(streamRequest, callback);
    latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
    Assert.assertEquals(status.get(), RestStatus.OK);

    List<SinglePartMIMEReaderCallbackImpl> singlePartMIMEReaderCallbacks =
        _mimeServerRequestHandler.getTestMultiPartMIMEReaderCallback().getSinglePartMIMEReaderCallbacks();
        _mimeServerRequestHandler.getTestMultiPartMIMEReaderCallback().getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), expectedParts.size());
    for (int i = 0; i < singlePartMIMEReaderCallbacks.size(); i++)
    {
      //Actual
      final SinglePartMIMEReaderCallbackImpl currentCallback = singlePartMIMEReaderCallbacks.get(i);
      //Expected
      final MIMEDataPart currentExpectedPart = expectedParts.get(i);

      Assert.assertEquals(currentCallback.getHeaders(), currentExpectedPart.getPartHeaders());
      Assert.assertEquals(currentCallback.getFinishedData(), currentExpectedPart.getPartData());
    }
  }

  private static class SinglePartMIMEReaderCallbackImpl implements SinglePartMIMEReaderCallback
  {
    final MultiPartMIMEReaderCallback _topLevelCallback;
    final MultiPartMIMEReader.SinglePartMIMEReader _singlePartMIMEReader;
    final ByteArrayOutputStream _byteArrayOutputStream = new ByteArrayOutputStream();
    Map<String, String> _headers;
    ByteString _finishedData = null;
    static int partCounter = 0;

    SinglePartMIMEReaderCallbackImpl(final MultiPartMIMEReaderCallback topLevelCallback,
                                     final MultiPartMIMEReader.SinglePartMIMEReader singlePartMIMEReader)
    {
      _topLevelCallback = topLevelCallback;
      _singlePartMIMEReader = singlePartMIMEReader;
      _headers = singlePartMIMEReader.dataSourceHeaders();
    }

    public Map<String, String> getHeaders()
    {
      return _headers;
    }

    public ByteString getFinishedData()
    {
      return _finishedData;
    }

    @Override
    public void onPartDataAvailable(ByteString partData)
    {
      try
      {
        _byteArrayOutputStream.write(partData.copyBytes());
      }
      catch (IOException ioException)
      {
        onStreamError(ioException);
      }
      _singlePartMIMEReader.requestPartData();
    }

    @Override
    public void onFinished()
    {
      partCounter++;
      _finishedData = ByteString.copy(_byteArrayOutputStream.toByteArray());
    }

    //Delegate to the top level for now for these two
    @Override
    public void onDrainComplete()
    {
      //This will end up failing the test.
      _topLevelCallback.onDrainComplete();
    }

    @Override
    public void onStreamError(Throwable throwable)
    {
      //MultiPartMIMEReader will end up calling onStreamError(e) on our top level callback
      //which will fail the test
    }
  }

  private static class MultiPartMIMEReaderCallbackImpl implements MultiPartMIMEReaderCallback
  {
    final Callback<StreamResponse> _r2callback;
    final List<SinglePartMIMEReaderCallbackImpl> _singlePartMIMEReaderCallbacks = new ArrayList<SinglePartMIMEReaderCallbackImpl>();

    MultiPartMIMEReaderCallbackImpl(final Callback<StreamResponse> r2callback)
    {
      _r2callback = r2callback;
    }

    public List<SinglePartMIMEReaderCallbackImpl> getSinglePartMIMEReaderCallbacks()
    {
      return _singlePartMIMEReaderCallbacks;
    }

    @Override
    public void onNewPart(MultiPartMIMEReader.SinglePartMIMEReader singlePartMIMEReader)
    {
      SinglePartMIMEReaderCallbackImpl singlePartMIMEReaderCallback = new SinglePartMIMEReaderCallbackImpl(this, singlePartMIMEReader);
      singlePartMIMEReader.registerReaderCallback(singlePartMIMEReaderCallback);
      _singlePartMIMEReaderCallbacks.add(singlePartMIMEReaderCallback);
      singlePartMIMEReader.requestPartData();
    }

    @Override
    public void onFinished()
    {
      RestResponse response = RestStatus.responseForStatus(RestStatus.OK, "");
      _r2callback.onSuccess(Messages.toStreamResponse(response));
    }

    @Override
    public void onDrainComplete()
    {
      RestException restException = new RestException(RestStatus.responseForStatus(406, "Not Acceptable"));
      _r2callback.onError(restException);
    }

    @Override
    public void onStreamError(Throwable throwable)
    {
      RestException restException = new RestException(RestStatus.responseForError(400, throwable));
      _r2callback.onError(restException);
    }
  }

  private static class MIMEServerRequestHandler implements StreamRequestHandler
  {
    private MultiPartMIMEReaderCallbackImpl _testMultiPartMIMEReaderCallback;

    MIMEServerRequestHandler()
    {
    }

    public MultiPartMIMEReaderCallbackImpl getTestMultiPartMIMEReaderCallback()
    {
      return _testMultiPartMIMEReaderCallback;
    }

    @Override
    public void handleRequest(StreamRequest request, RequestContext requestContext,
        final Callback<StreamResponse> callback)
    {
      try
      {
        MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(request);
        _testMultiPartMIMEReaderCallback = new MultiPartMIMEReaderCallbackImpl(callback);
        reader.registerReaderCallback(_testMultiPartMIMEReaderCallback);
      }
      catch (MultiPartIllegalFormatException illegalMimeFormatException)
      {
        RestException restException = new RestException(RestStatus.responseForError(400, illegalMimeFormatException));
        callback.onError(restException);
      }
    }
  }
}