Examples with TestingResponse com.facebook.airlift.http.client.testing.TestingResponse used on opensource projects

Example 6 with TestingResponse

use of com.facebook.airlift.http.client.testing.TestingResponse in project presto by prestodb.

the class MockExchangeRequestProcessor method handle.

@Override
public Response handle(Request request) {
    if (request.getMethod().equalsIgnoreCase("DELETE")) {
        return new TestingResponse(HttpStatus.NO_CONTENT, ImmutableListMultimap.of(), new byte[0]);
    }
    // verify we got a data size and it parses correctly
    assertTrue(!request.getHeaders().get(PrestoHeaders.PRESTO_MAX_SIZE).isEmpty());
    DataSize maxSize = DataSize.valueOf(request.getHeader(PrestoHeaders.PRESTO_MAX_SIZE));
    assertTrue(maxSize.compareTo(expectedMaxSize) <= 0);
    requestMaxSizes.add(maxSize);
    RequestLocation requestLocation = new RequestLocation(request.getUri());
    URI location = requestLocation.getLocation();
    BufferResult result = buffers.getUnchecked(location).getPages(requestLocation.getSequenceId(), maxSize);
    byte[] bytes = new byte[0];
    HttpStatus status;
    if (!result.getSerializedPages().isEmpty()) {
        DynamicSliceOutput sliceOutput = new DynamicSliceOutput(64);
        PagesSerdeUtil.writeSerializedPages(sliceOutput, result.getSerializedPages());
        bytes = dataChanger.apply(sliceOutput.slice().getBytes());
        status = HttpStatus.OK;
    } else {
        status = HttpStatus.NO_CONTENT;
    }
    return new TestingResponse(status, ImmutableListMultimap.of(CONTENT_TYPE, PRESTO_PAGES, PRESTO_TASK_INSTANCE_ID, String.valueOf(result.getTaskInstanceId()), PRESTO_PAGE_TOKEN, String.valueOf(result.getToken()), PRESTO_PAGE_NEXT_TOKEN, String.valueOf(result.getNextToken()), PRESTO_BUFFER_COMPLETE, String.valueOf(result.isBufferComplete())), bytes);
}

Example 7 with TestingResponse

use of com.facebook.airlift.http.client.testing.TestingResponse in project presto by prestodb.

the class TestingExchangeHttpClientHandler method handle.

@Override
public Response handle(Request request) {
    ImmutableList<String> parts = ImmutableList.copyOf(Splitter.on("/").omitEmptyStrings().split(request.getUri().getPath()));
    if (request.getMethod().equals("DELETE")) {
        assertEquals(parts.size(), 1);
        return new TestingResponse(HttpStatus.OK, ImmutableListMultimap.of(), new byte[0]);
    }
    assertEquals(parts.size(), 2);
    String taskId = parts.get(0);
    int pageToken = Integer.parseInt(parts.get(1));
    ImmutableListMultimap.Builder<String, String> headers = ImmutableListMultimap.builder();
    headers.put(PRESTO_TASK_INSTANCE_ID, "task-instance-id");
    headers.put(PRESTO_PAGE_TOKEN, String.valueOf(pageToken));
    TestingTaskBuffer taskBuffer = taskBuffers.getUnchecked(taskId);
    Page page = taskBuffer.getPage(pageToken);
    headers.put(CONTENT_TYPE, PRESTO_PAGES);
    if (page != null) {
        headers.put(PRESTO_PAGE_NEXT_TOKEN, String.valueOf(pageToken + 1));
        headers.put(PRESTO_BUFFER_COMPLETE, String.valueOf(false));
        DynamicSliceOutput output = new DynamicSliceOutput(256);
        PagesSerdeUtil.writePages(PAGES_SERDE, output, page);
        return new TestingResponse(HttpStatus.OK, headers.build(), output.slice().getInput());
    } else if (taskBuffer.isFinished()) {
        headers.put(PRESTO_PAGE_NEXT_TOKEN, String.valueOf(pageToken));
        headers.put(PRESTO_BUFFER_COMPLETE, String.valueOf(true));
        return new TestingResponse(HttpStatus.OK, headers.build(), new byte[0]);
    } else {
        headers.put(PRESTO_PAGE_NEXT_TOKEN, String.valueOf(pageToken));
        headers.put(PRESTO_BUFFER_COMPLETE, String.valueOf(false));
        return new TestingResponse(HttpStatus.NO_CONTENT, headers.build(), new byte[0]);
    }
}