com.linkedin.pinot.server.request.ScheduledRequestHandlerTest.java Source code

Java tutorial

Introduction

Here is the source code for com.linkedin.pinot.server.request.ScheduledRequestHandlerTest.java

Source

/**
 * Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com)
 *
 * 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.pinot.server.request;

import com.google.common.util.concurrent.ListenableFuture;
import com.linkedin.pinot.common.data.FieldSpec;
import com.linkedin.pinot.common.exception.QueryException;
import com.linkedin.pinot.common.metrics.ServerMetrics;
import com.linkedin.pinot.common.query.QueryExecutor;
import com.linkedin.pinot.common.query.QueryRequest;
import com.linkedin.pinot.common.request.BrokerRequest;
import com.linkedin.pinot.common.request.InstanceRequest;
import com.linkedin.pinot.common.utils.DataSchema;
import com.linkedin.pinot.common.utils.DataTable;
import com.linkedin.pinot.core.common.datatable.DataTableBuilder;
import com.linkedin.pinot.core.common.datatable.DataTableFactory;
import com.linkedin.pinot.core.common.datatable.DataTableImplV2;
import com.linkedin.pinot.core.query.executor.ServerQueryExecutorV1Impl;
import com.linkedin.pinot.core.query.scheduler.QueryScheduler;
import com.linkedin.pinot.serde.SerDe;
import com.yammer.metrics.core.MetricsRegistry;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.Arrays;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.apache.thrift.protocol.TCompactProtocol;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.BeforeTest;
import org.testng.annotations.Test;

import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class ScheduledRequestHandlerTest {

    public static Logger LOGGER = LoggerFactory.getLogger(ScheduledRequestHandlerTest.class);

    private ServerMetrics serverMetrics;
    private ChannelHandlerContext channelHandlerContext;
    private QueryScheduler queryScheduler;
    private QueryExecutor queryExecutor;

    @BeforeTest
    public void setupTestMethod() {
        serverMetrics = new ServerMetrics(new MetricsRegistry());
        channelHandlerContext = mock(ChannelHandlerContext.class, RETURNS_DEEP_STUBS);
        when(channelHandlerContext.channel().remoteAddress()).thenAnswer(new Answer<InetSocketAddress>() {
            @Override
            public InetSocketAddress answer(InvocationOnMock invocationOnMock) throws Throwable {
                return new InetSocketAddress("localhost", 60000);
            }
        });

        queryScheduler = mock(QueryScheduler.class);
        queryExecutor = new ServerQueryExecutorV1Impl();
    }

    @Test
    public void testBadRequest() throws Exception {
        ScheduledRequestHandler handler = new ScheduledRequestHandler(queryScheduler, serverMetrics);
        String requestBadString = "foobar";
        byte[] requestData = requestBadString.getBytes();
        ByteBuf buffer = Unpooled.wrappedBuffer(requestData);
        ListenableFuture<byte[]> response = handler.processRequest(channelHandlerContext, buffer);
        // The handler method is expected to return immediately
        Assert.assertTrue(response.isDone());
        byte[] responseBytes = response.get();
        Assert.assertTrue(responseBytes.length > 0);
        DataTable expectedDT = new DataTableImplV2();
        expectedDT.addException(QueryException.INTERNAL_ERROR);
        Assert.assertEquals(responseBytes, expectedDT.toBytes());
    }

    private InstanceRequest getInstanceRequest() {
        InstanceRequest request = new InstanceRequest();
        request.setBrokerId("broker");
        request.setEnableTrace(false);
        request.setRequestId(1);
        request.setSearchSegments(Arrays.asList("segment1", "segment2"));
        request.setQuery(new BrokerRequest());
        return request;
    }

    private ByteBuf getSerializedInstanceRequest(InstanceRequest request) {
        SerDe serDe = new SerDe(new TCompactProtocol.Factory());
        byte[] requestData = serDe.serialize(request);
        return Unpooled.wrappedBuffer(requestData);
    }

    @Test
    public void testQueryProcessingException() throws Exception {
        ScheduledRequestHandler handler = new ScheduledRequestHandler(new QueryScheduler(queryExecutor) {
            @Override
            public ListenableFuture<DataTable> submit(QueryRequest queryRequest) {
                return queryWorkers.submit(new Callable<DataTable>() {
                    @Override
                    public DataTable call() throws Exception {
                        throw new RuntimeException("query processing error");
                    }
                });
            }
        }, serverMetrics);

        ByteBuf requestBuf = getSerializedInstanceRequest(getInstanceRequest());
        ListenableFuture<byte[]> responseFuture = handler.processRequest(channelHandlerContext, requestBuf);
        byte[] bytes = responseFuture.get(2, TimeUnit.SECONDS);
        // we get DataTable with exception information in case of query processing exception
        Assert.assertTrue(bytes.length > 0);
        DataTable expectedDT = new DataTableImplV2();
        expectedDT.addException(QueryException.INTERNAL_ERROR);
        Assert.assertEquals(bytes, expectedDT.toBytes());
    }

    @Test
    public void testValidQueryResponse()
            throws InterruptedException, ExecutionException, TimeoutException, IOException {
        ScheduledRequestHandler handler = new ScheduledRequestHandler(new QueryScheduler(queryExecutor) {
            @Override
            public ListenableFuture<DataTable> submit(QueryRequest queryRequest) {
                return queryRunners.submit(new Callable<DataTable>() {
                    @Override
                    public DataTable call() throws Exception {
                        String[] columns = new String[] { "foo", "bar" };
                        FieldSpec.DataType[] columnTypes = new FieldSpec.DataType[] { FieldSpec.DataType.STRING,
                                FieldSpec.DataType.INT };
                        DataSchema dataSchema = new DataSchema(columns, columnTypes);
                        DataTableBuilder dtBuilder = new DataTableBuilder(dataSchema);
                        dtBuilder.startRow();
                        dtBuilder.setColumn(0, "mars");
                        dtBuilder.setColumn(1, 10);
                        dtBuilder.finishRow();
                        dtBuilder.startRow();
                        dtBuilder.setColumn(0, "jupiter");
                        dtBuilder.setColumn(1, 100);
                        dtBuilder.finishRow();
                        return dtBuilder.build();
                    }
                });
            }
        }, serverMetrics);

        ByteBuf requestBuf = getSerializedInstanceRequest(getInstanceRequest());
        ListenableFuture<byte[]> responseFuture = handler.processRequest(channelHandlerContext, requestBuf);
        byte[] responseBytes = responseFuture.get(2, TimeUnit.SECONDS);
        DataTable responseDT = DataTableFactory.getDataTable(responseBytes);
        Assert.assertEquals(responseDT.getNumberOfRows(), 2);
        Assert.assertEquals(responseDT.getString(0, 0), "mars");
        Assert.assertEquals(responseDT.getInt(0, 1), 10);
        Assert.assertEquals(responseDT.getString(1, 0), "jupiter");
        Assert.assertEquals(responseDT.getInt(1, 1), 100);
    }

}