com.linkedin.pinot.common.http.MultiGetRequestTest.java Source code

Introduction

Here is the source code for com.linkedin.pinot.common.http.MultiGetRequestTest.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.common.http;

import com.google.common.base.Throwables;
import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.SocketTimeoutException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import org.apache.commons.httpclient.MultiThreadedHttpConnectionManager;
import org.apache.commons.httpclient.methods.GetMethod;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.AfterTest;
import org.testng.annotations.BeforeTest;
import org.testng.annotations.Test;

public class MultiGetRequestTest {
    private static final Logger LOGGER = LoggerFactory.getLogger(MultiGetRequest.class);

    List<HttpServer> servers = new ArrayList<>();
    final int portStart = 14000;
    final String SUCCESS_MSG = "success";
    final String ERROR_MSG = "error";
    final String TIMEOUT_MSG = "Timeout";
    final int SUCCESS_CODE = 200;
    final int ERROR_CODE = 500;
    final String URI_PATH = "/foo";
    final int TIMEOUT_MS = 5000;

    @BeforeTest
    public void setUpTest() throws IOException {

        startServer(portStart, createHandler(SUCCESS_CODE, SUCCESS_MSG, 0));
        startServer(portStart + 1, createHandler(ERROR_CODE, ERROR_MSG, 0));
        startServer(portStart + 2, createHandler(SUCCESS_CODE, TIMEOUT_MSG, TIMEOUT_MS));

    }

    @AfterTest
    public void tearDownTest() {
        for (HttpServer server : servers) {
            server.stop(0);
        }
    }

    private HttpHandler createHandler(final int status, final String msg, final int sleepTimeMs) {
        return new HttpHandler() {
            @Override
            public void handle(HttpExchange httpExchange) throws IOException {
                if (sleepTimeMs > 0) {
                    try {
                        Thread.sleep(sleepTimeMs);
                    } catch (InterruptedException e) {
                        LOGGER.info("Handler interrupted during sleep");
                    }
                }
                httpExchange.sendResponseHeaders(status, msg.length());
                OutputStream responseBody = httpExchange.getResponseBody();
                responseBody.write(msg.getBytes());
                responseBody.close();
            }
        };
    }

    private HttpServer startServer(int port, HttpHandler handler) throws IOException {
        final HttpServer server = HttpServer.create(new InetSocketAddress(port), 0);
        server.createContext(URI_PATH, handler);
        new Thread(new Runnable() {
            @Override
            public void run() {
                server.start();
            }
        }).start();
        servers.add(server);
        return server;
    }

    @Test
    public void testMultiGet() {
        MultiGetRequest mget = new MultiGetRequest(Executors.newCachedThreadPool(),
                new MultiThreadedHttpConnectionManager());
        List<String> urls = Arrays.asList("http://localhost:" + String.valueOf(portStart) + URI_PATH,
                "http://localhost:" + String.valueOf(portStart + 1) + URI_PATH,
                "http://localhost:" + String.valueOf(portStart + 2) + URI_PATH,
                // 2nd request to the same server
                "http://localhost:" + String.valueOf(portStart) + URI_PATH);
        // timeout value needs to be less than 5000ms set above for
        // third server
        final int requestTimeoutMs = 1000;
        CompletionService<GetMethod> completionService = mget.execute(urls, requestTimeoutMs);
        int success = 0;
        int errors = 0;
        int timeouts = 0;
        for (int i = 0; i < urls.size(); i++) {
            GetMethod getMethod = null;
            try {
                getMethod = completionService.take().get();
                if (getMethod.getStatusCode() >= 300) {
                    ++errors;
                    Assert.assertEquals(getMethod.getResponseBodyAsString(), ERROR_MSG);
                } else {
                    ++success;
                    Assert.assertEquals(getMethod.getResponseBodyAsString(), SUCCESS_MSG);
                }

            } catch (InterruptedException e) {
                LOGGER.error("Interrupted", e);
                ++errors;
            } catch (ExecutionException e) {
                if (Throwables.getRootCause(e) instanceof SocketTimeoutException) {
                    LOGGER.debug("Timeout");
                    ++timeouts;
                } else {
                    LOGGER.error("Error", e);
                    ++errors;
                }
            } catch (IOException e) {
                ++errors;
            }
        }

        Assert.assertEquals(2, success);
        Assert.assertEquals(1, errors);
        Assert.assertEquals(1, timeouts);
    }

}