alluxio.worker.grpc.AbstractWriteHandlerTest.java Source code

Java tutorial

Introduction

Here is the source code for alluxio.worker.grpc.AbstractWriteHandlerTest.java

Source

/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.worker.grpc;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;

import alluxio.Constants;
import alluxio.grpc.Chunk;
import alluxio.grpc.RequestType;
import alluxio.grpc.WriteRequestCommand;
import alluxio.grpc.WriteResponse;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.network.protocol.databuffer.ByteArrayDataBuffer;
import alluxio.security.authentication.AuthenticatedUserInfo;
import alluxio.util.CommonUtils;
import alluxio.util.WaitForOptions;
import alluxio.util.io.BufferUtils;

import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.StatusException;
import io.grpc.stub.StreamObserver;
import io.netty.buffer.ByteBuf;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.rules.TemporaryFolder;
import org.mockito.ArgumentCaptor;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.TimeoutException;

/**
 * Unit tests for {@link AbstractWriteHandler}.
 */
public abstract class AbstractWriteHandlerTest {
    private static final Random RANDOM = new Random();
    protected static final int CHUNK_SIZE = 1024;
    protected static final long TEST_BLOCK_ID = 1L;
    protected static final long TEST_MOUNT_ID = 10L;
    protected AbstractWriteHandler mWriteHandler;
    protected StreamObserver<WriteResponse> mResponseObserver;
    protected AuthenticatedUserInfo mUserInfo = new AuthenticatedUserInfo();
    protected List<WriteResponse> mResponses = new ArrayList<>();
    protected boolean mResponseCompleted;
    protected Throwable mError;

    @Rule
    public ExpectedException mExpectedException = ExpectedException.none();

    @Rule
    public TemporaryFolder mTestFolder = new TemporaryFolder();

    @Test
    public void writeEmptyFile() throws Exception {
        mWriteHandler.write(newWriteRequestCommand(0));
        mWriteHandler.onCompleted();
        waitForResponses();
        checkComplete(mResponseObserver);
    }

    @Test
    public void writeNonEmptyFile() throws Exception {
        long len = 0;
        long checksum = 0;
        mWriteHandler.write(newWriteRequestCommand(0));
        for (int i = 0; i < 128; i++) {
            DataBuffer dataBuffer = newDataBuffer(CHUNK_SIZE);
            checksum += getChecksum(dataBuffer);
            mWriteHandler.write(newWriteRequest(dataBuffer));
            len += CHUNK_SIZE;
        }
        // EOF.
        mWriteHandler.onCompleted();
        waitForResponses();
        checkComplete(mResponseObserver);
        checkWriteData(checksum, len);
    }

    @Test
    public void cancel() throws Exception {
        long len = 0;
        long checksum = 0;
        mWriteHandler.write(newWriteRequestCommand(0));
        for (int i = 0; i < 1; i++) {
            DataBuffer dataBuffer = newDataBuffer(CHUNK_SIZE);
            checksum += getChecksum(dataBuffer);
            mWriteHandler.write(newWriteRequest(dataBuffer));
            len += CHUNK_SIZE;
        }
        // Cancel.
        mWriteHandler.onCancel();

        waitForResponses();
        checkComplete(mResponseObserver);
        // Our current implementation does not really abort the file when the write is cancelled.
        // The client issues another request to block worker to abort it.
        checkWriteData(checksum, len);
    }

    @Test
    public void cancelIgnoreError() throws Exception {
        long len = 0;
        long checksum = 0;
        mWriteHandler.write(newWriteRequestCommand(0));
        for (int i = 0; i < 1; i++) {
            DataBuffer dataBuffer = newDataBuffer(CHUNK_SIZE);
            checksum += getChecksum(dataBuffer);
            mWriteHandler.write(newWriteRequest(dataBuffer));
            len += CHUNK_SIZE;
        }
        // Cancel.
        mWriteHandler.onCancel();
        mWriteHandler.onError(Status.CANCELLED.asRuntimeException());

        waitForResponses();
        checkComplete(mResponseObserver);
        checkWriteData(checksum, len);
        verify(mResponseObserver, never()).onError(any(Throwable.class));
    }

    @Test
    public void writeInvalidOffsetFirstRequest() throws Exception {
        // The write request contains an invalid offset
        mWriteHandler.write(newWriteRequestCommand(1));
        waitForResponses();
        checkErrorCode(mResponseObserver, Status.Code.INVALID_ARGUMENT);
    }

    @Test
    public void writeInvalidOffsetLaterRequest() throws Exception {
        mWriteHandler.write(newWriteRequestCommand(0));
        // The write request contains an invalid offset
        mWriteHandler.write(newWriteRequestCommand(1));
        waitForResponses();
        checkErrorCode(mResponseObserver, Status.Code.INVALID_ARGUMENT);
    }

    @Test
    public void ErrorReceived() throws Exception {
        mWriteHandler.onError(new IOException("test exception"));
    }

    @Test
    public void ErrorReceivedAfterRequest() throws Exception {
        mWriteHandler.write(newWriteRequestCommand(0));
        mWriteHandler.onCompleted();
        mWriteHandler.onError(new IOException("test exception"));
    }

    /**
     * Checks the file content matches expectation (file length and file checksum).
     *
     * @param expectedChecksum the expected checksum of the file
     * @param size the file size in bytes
     */
    protected void checkWriteData(long expectedChecksum, long size) throws IOException {
        long actualChecksum = 0;
        long actualSize = 0;

        byte[] buffer = new byte[(int) Math.min(Constants.KB, size)];
        try (InputStream input = getWriteDataStream()) {
            int bytesRead;
            while ((bytesRead = input.read(buffer)) >= 0) {
                for (int i = 0; i < bytesRead; i++) {
                    actualChecksum += BufferUtils.byteToInt(buffer[i]);
                    actualSize++;
                }
            }
        }

        assertEquals(expectedChecksum, actualChecksum);
        assertEquals(size, actualSize);
    }

    /**
     * Waits for an error.
     *
     * @return the error
     */
    protected Throwable getError(final StreamObserver<WriteResponse> responseObserver) {
        ArgumentCaptor<Throwable> captor = ArgumentCaptor.forClass(Throwable.class);
        verify(responseObserver).onError(captor.capture());
        return captor.getValue();
    }

    /**
     * Checks an error is returned with given code.
     *
     * @param responseObserver the response stream observer
     * @param code the expected error code
     */
    protected void checkErrorCode(final StreamObserver<WriteResponse> responseObserver, Status.Code code) {
        Throwable t = getError(responseObserver);
        assertTrue(t instanceof StatusException);
        assertEquals(code, ((StatusException) t).getStatus().getCode());
    }

    /**
     * Checks that the response is completed.
     */
    protected void checkComplete(final StreamObserver<WriteResponse> responseObserver) {
        verify(responseObserver).onCompleted();
    }

    protected alluxio.grpc.WriteRequest newWriteRequestCommand(long offset) {
        return alluxio.grpc.WriteRequest.newBuilder().setCommand(WriteRequestCommand.newBuilder()
                .setId(TEST_BLOCK_ID).setOffset(offset).setType(getWriteRequestType())).build();
    }

    /**
     * Builds the write request with data.
     *
     * @param buffer the data to write
     * @return the write request
     */
    protected alluxio.grpc.WriteRequest newWriteRequest(DataBuffer buffer) {
        alluxio.grpc.WriteRequest writeRequest = alluxio.grpc.WriteRequest.newBuilder()
                .setChunk(Chunk.newBuilder().setData(ByteString.copyFrom(buffer.getReadOnlyByteBuffer()))).build();
        return writeRequest;
    }

    /**
     * @return the write type of the request
     */
    protected abstract RequestType getWriteRequestType();

    /**
     * @return the data written as an InputStream
     */
    protected abstract InputStream getWriteDataStream() throws IOException;

    /**
     * @param len length of the data buffer
     * @return a newly created data buffer
     */
    protected DataBuffer newDataBuffer(int len) {
        byte[] buf = new byte[len];
        for (int i = 0; i < len; i++) {
            byte value = (byte) (RANDOM.nextInt() % Byte.MAX_VALUE);
            buf[i] = value;
        }
        return new ByteArrayDataBuffer(buf, 0, len);
    }

    /**
     * @param buffer buffer to get checksum
     * @return the checksum
     */
    public static long getChecksum(DataBuffer buffer) {
        return getChecksum((ByteBuf) buffer.getNettyOutput());
    }

    /**
     * @param buffer buffer to get checksum
     * @return the checksum
     */
    public static long getChecksum(ByteBuf buffer) {
        long ret = 0;
        for (int i = 0; i < buffer.capacity(); i++) {
            ret += BufferUtils.byteToInt(buffer.getByte(i));
        }
        return ret;
    }

    /**
     * Waits for response messages.
     */
    protected void waitForResponses() throws TimeoutException, InterruptedException {
        CommonUtils.waitFor("response", () -> mResponseCompleted || mError != null,
                WaitForOptions.defaults().setTimeoutMs(Constants.MINUTE_MS));
    }

    protected void setupResponseTrigger() {
        doAnswer(args -> {
            mResponseCompleted = true;
            return null;
        }).when(mResponseObserver).onCompleted();
        doAnswer(args -> {
            mResponseCompleted = true;
            mError = args.getArgumentAt(0, Throwable.class);
            return null;
        }).when(mResponseObserver).onError(any(Throwable.class));
        doAnswer((args) -> {
            // make a copy of response data before it is released
            mResponses.add(WriteResponse.parseFrom(args.getArgumentAt(0, WriteResponse.class).toByteString()));
            return null;
        }).when(mResponseObserver).onNext(any(WriteResponse.class));
    }
}