Examples with VariableByteStringViewer com.linkedin.multipart.utils.VariableByteStringViewer used on opensource projects

Example 1 with VariableByteStringViewer

use of com.linkedin.multipart.utils.VariableByteStringViewer in project rest.li by linkedin.

the class AbstractMIMEUnitTest method mockR2AndWrite.

// This is used when we need to mock out R2 and write a payload for our reader to read
protected void mockR2AndWrite(final ByteString payload, final int chunkSize, final String contentType) {
    _entityStream = mock(EntityStream.class);
    _readHandle = mock(ReadHandle.class);
    _streamRequest = mock(StreamRequest.class);
    // We have to use the AtomicReference holder technique to modify the current remaining buffer since the inner class
    // in doAnswer() can only access final variables.
    final AtomicReference<MultiPartMIMEReader.R2MultiPartMIMEReader> r2Reader = new AtomicReference<>();
    // This takes the place of VariableByteStringWriter if we were to use R2 directly.
    final VariableByteStringViewer variableByteStringViewer = new VariableByteStringViewer(payload, chunkSize);
    // When data is requested, we write
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            final MultiPartMIMEReader.R2MultiPartMIMEReader reader = r2Reader.get();
            Object[] args = invocation.getArguments();
            // will always be 1 since MultiPartMIMEReader only does _rh.request(1)
            final int chunksRequested = (Integer) args[0];
            for (int i = 0; i < chunksRequested; i++) {
                // Our tests will run into a stack overflow unless we use a thread pool here to fire off the callbacks.
                // Especially in cases where the chunk size is 1. When the chunk size is one, the MultiPartMIMEReader
                // ends up doing many _rh.request(1) since each write is only 1 byte.
                // R2 uses a different technique to avoid stack overflows here which is unnecessary to emulate.
                _scheduledExecutorService.submit(new Runnable() {

                    @Override
                    public void run() {
                        ByteString clientData = variableByteStringViewer.onWritePossible();
                        if (clientData.equals(ByteString.empty())) {
                            reader.onDone();
                        } else {
                            reader.onDataAvailable(clientData);
                        }
                    }
                });
            }
            return null;
        }
    }).when(_readHandle).request(isA(Integer.class));
    // We need a final version of the read handle since its passed to an inner class below.
    final ReadHandle readHandleRef = _readHandle;
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            final MultiPartMIMEReader.R2MultiPartMIMEReader reader = (MultiPartMIMEReader.R2MultiPartMIMEReader) args[0];
            r2Reader.set(reader);
            // R2 calls init immediately upon setting the reader
            reader.onInit(readHandleRef);
            return null;
        }
    }).when(_entityStream).setReader(isA(MultiPartMIMEReader.R2MultiPartMIMEReader.class));
    when(_streamRequest.getEntityStream()).thenReturn(_entityStream);
    final String contentTypeHeader = contentType + ";somecustomparameter=somecustomvalue" + ";anothercustomparameter=anothercustomvalue";
    when(_streamRequest.getHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER)).thenReturn(contentTypeHeader);
}