Start line:  
End line:  

Snippet Preview

Snippet HTML Code

Stack Overflow Questions
  /*
   * 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.facebook.presto.server;
 
 
 import java.net.URI;
 import java.util.List;
 
 import static com.google.common.base.Charsets.UTF_8;
 import static io.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler;
 import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
 import static io.airlift.http.client.JsonResponseHandler.createJsonResponseHandler;
 import static io.airlift.http.client.Request.Builder.prepareDelete;
 import static io.airlift.http.client.Request.Builder.prepareGet;
 import static io.airlift.http.client.Request.Builder.preparePost;
 import static io.airlift.http.client.ResponseHandlerUtils.propagate;
 import static io.airlift.http.client.StaticBodyGenerator.createStaticBodyGenerator;
 import static io.airlift.http.client.StatusResponseHandler.createStatusResponseHandler;
 import static io.airlift.json.JsonCodec.jsonCodec;
 import static org.testng.Assert.assertEquals;
 
 @Test(singleThreaded = true)
 public class TestQueryResourceServer
 {
     private HttpClient client;
     private TestingHttpServer server;
 
     @BeforeMethod
     public void setup()
             throws Exception
     {
         Bootstrap app = new Bootstrap(
                 new TestingNodeModule(),
                 new InMemoryEventModule(),
                 new TestingHttpServerModule(),
                 new JsonModule(),
                 new JaxrsModule(),
                 new Module()
                 {
                     @Override
                     public void configure(Binder binder)
                     {
                         binder.bind(QueryResource.class).in(.);
                         binder.bind(StageResource.class).in(.);
                         binder.bind(TaskResource.class).in(.);
                         binder.bind(QueryManager.class).to(MockQueryManager.class).in(.);
                         binder.bind(MockTaskManager.class).in(.);
                         binder.bind(TaskManager.class).to(Key.get(MockTaskManager.class)).in(.);
                         binder.bind(PagesMapper.class).in(.);
                         binder.bind(InternalNodeManager.class).to(InMemoryNodeManager.class).in(.);
                        binder.bind(NodeManager.class).to(Key.get(InternalNodeManager.class)).in(.);
                        binder.bind(LocationFactory.class).to(HttpLocationFactory.class).in(.);
                    }
                });
        Injector injector = app
                .strictConfig()
                .doNotInitializeLogging()
                .initialize();
         = injector.getInstance(LifeCycleManager.class);
         = injector.getInstance(TestingHttpServer.class);
         = new JettyHttpClient();
    }
    public void teardown()
            throws Exception
    {
        if ( != null) {
            .stop();
        }
        if ( != null) {
            .close();
        }
    }
    @Test(enabled = false)
    public void testQuery()
            throws Exception
    {
        URI location = .execute(preparePost().setUri(uriFor("/v1/query")).setBodyGenerator(createStaticBodyGenerator("query")).build(), new CreatedResponseHandler());
        assertQueryStatus(location.);
        QueryInfo queryInfo = .execute(prepareGet().setUri(location).build(), createJsonResponseHandler(jsonCodec(QueryInfo.class)));
        TaskInfo taskInfo = queryInfo.getOutputStage().getTasks().get(0);
        URI outputLocation = uriFor("/v1/task/" + taskInfo.getTaskId() + "/results/out");
        long sequenceId = 0;
        PagesResponse response = .execute(
                prepareGet().setUri(uriBuilderFrom(outputLocation).appendPath(String.valueOf(sequenceId)).build()).build(),
                new PageResponseHandler());
        List<Pagepages = response.getPages();
        assertEquals(countPositions(pages), 220);
        assertQueryStatus(location.);
        sequenceId += pages.size();
        response = .execute(
                prepareGet().setUri(uriBuilderFrom(outputLocation).appendPath(String.valueOf(sequenceId)).build()).build(),
                new PageResponseHandler());
        pages = response.getPages();
        assertEquals(countPositions(pages), 44 + 48);
        sequenceId += pages.size();
        response = .execute(
                prepareGet().setUri(uriBuilderFrom(outputLocation).appendPath(String.valueOf(sequenceId)).build()).build(),
                new PageResponseHandler());
        pages = response.getPages();
        assertEquals(countPositions(pages), 0);
        assertQueryStatus(location.);
        // cancel the query
        StatusResponse cancelResponse = .execute(prepareDelete().setUri(location).build(), createStatusResponseHandler());
        assertQueryStatus(location.);
        assertEquals(cancelResponse.getStatusCode(), ..code());
    }
    private int countPositions(List<Pagepages)
    {
        int count = 0;
        for (Page page : pages) {
            count += page.getPositionCount();
        }
        return count;
    }
    private void assertQueryStatus(URI locationQueryState expectedQueryState)
    {
        URI statusUri = uriBuilderFrom(location).build();
        if (expectedQueryState == . && response.getStatusCode() == ..code()) {
            // when query finishes the server may delete it
            return;
        }
        QueryInfo queryInfo = response.getValue();
        assertEquals(queryInfo.getState(), expectedQueryState);
    }
    private URI uriFor(String path)
    {
        return .getBaseUrl().resolve(path);
    }
    private static class CreatedResponseHandler
            implements ResponseHandler<URIRuntimeException>
    {
        @Override
        public URI handleException(Request requestException exception)
        {
            throw propagate(requestexception);
        }
        @Override
        public URI handle(Request requestResponse response)
        {
            if (response.getStatusCode() != ..code()) {
                throw new UnexpectedResponseException(
                        String.format("Expected response code to be 201 CREATED, but was %s %s"response.getStatusCode(), response.getStatusMessage()),
                        request,
                        response);
            }
            String location = response.getHeader("Location");
            if (location == null) {
                throw new UnexpectedResponseException("Response does not contain a Location header"requestresponse);
            }
            return URI.create(location);
        }
    }
New to GrepCode? Check out our FAQ X