Examples with IndividualBody com.linkedin.restli.common.multiplexer.IndividualBody used on opensource projects

Example 6 with IndividualBody

use of com.linkedin.restli.common.multiplexer.IndividualBody in project rest.li by linkedin.

the class TestMultiplexedRequestHandlerImpl method testIndividualRequestInheritHeadersAndCookiesFromEnvelopeRequest.

@Test(dataProvider = "multiplexerConfigurations")
public void testIndividualRequestInheritHeadersAndCookiesFromEnvelopeRequest(MultiplexerRunMode multiplexerRunMode) throws Exception {
    // When some request headers/cookies are passed in the envelope, we need to ensure
    // they are properly included in each of the individual requests sent to
    // MultiplexerSingletonFilter and request handler. In the high level, the expected behavior should be:
    // 1. IndividualRequest that is passed to MultiplexerSingletonFilter should already have headers inherited from the
    //    envelope request.
    // 2. RestRequest that is passed to the request handler should have both headers and cookies inherited from the
    //    envelope request.
    // Create a mockHandler. Captures all headers and cookies found in the request.
    final Map<String, String> headers = new TreeMap<String, String>(String.CASE_INSENSITIVE_ORDER);
    final Map<String, String> cookies = new HashMap<String, String>();
    SynchronousRequestHandler mockHandler = new SynchronousRequestHandler() {

        @Override
        public RestResponse handleRequestSync(RestRequest request, RequestContext requestContext) {
            try {
                headers.putAll(request.getHeaders());
                for (HttpCookie cookie : CookieUtil.decodeCookies(request.getCookies())) {
                    cookies.put(cookie.getName(), cookie.getValue());
                }
                return fakeIndRestResponse(jsonBodyToByteString(new IndividualBody()));
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    };
    // Create a mock MultiplexerSingletonFilter to put request headers inside another headers so
    // we can do assertion on it later.
    final Map<String, String> headersSeenInMuxFilter = new TreeMap<String, String>(String.CASE_INSENSITIVE_ORDER);
    MultiplexerSingletonFilter muxFilterWithSimulatedFailures = new MultiplexerSingletonFilter() {

        @Override
        public IndividualRequest filterIndividualRequest(IndividualRequest request) {
            headersSeenInMuxFilter.putAll(request.getHeaders());
            return request;
        }

        @Override
        public IndividualResponse filterIndividualResponse(IndividualResponse response) {
            return response;
        }
    };
    // Prepare request to mux handler
    FutureCallback<RestResponse> callback = new FutureCallback<RestResponse>();
    RequestContext requestContext = new RequestContext();
    Map<String, IndividualRequest> individualRequests = ImmutableMap.of("0", fakeIndRequest("/request", ImmutableMap.of("X-IndividualHeader", "individualHeader", "X-OverridableHeader", "overrideHeader"), Collections.<String, IndividualRequest>emptyMap()));
    Set<String> headerWhiteList = new HashSet<String>();
    headerWhiteList.add("X-IndividualHeader");
    headerWhiteList.add("X-OverridableHeader");
    // Prepare mux request with cookie
    RestRequest muxRequest = new RestRequestBuilder(fakeMuxRestRequest(individualRequests)).addCookie("cookie1=cookie1Value; cookie2=cookie2Value").addHeaderValue("X-overridableheader", "originalHeader").addHeaderValue("X-Envelope", "envelopeHeaderValue").build();
    // Create mux handler instance
    MultiplexedRequestHandlerImpl multiplexer = createMultiplexer(mockHandler, muxFilterWithSimulatedFailures, headerWhiteList, individualRequests.size(), multiplexerRunMode);
    try {
        multiplexer.handleRequest(muxRequest, requestContext, callback);
    } catch (Exception e) {
        fail("Multiplexer should not throw exception", e);
    }
    RestResponse muxRestResponse = callback.get();
    // Assert multiplexed request should return a 200 status code
    assertEquals(muxRestResponse.getStatus(), 200, "Multiplexer should return 200");
    MultiplexedResponseContent muxResponseContent = new MultiplexedResponseContent(DataMapConverter.bytesToDataMap(muxRestResponse.getHeaders(), muxRestResponse.getEntity()));
    IndividualResponse response = muxResponseContent.getResponses().get("0");
    assertEquals(response.getStatus().intValue(), 200, "Individual request should not fail. Response body is: " + response.getBody().toString());
    Map<String, String> expectedHeaders = new TreeMap<String, String>(String.CASE_INSENSITIVE_ORDER);
    expectedHeaders.putAll(ImmutableMap.of(RestConstants.HEADER_CONTENT_TYPE, RestConstants.HEADER_VALUE_APPLICATION_JSON, "X-IndividualHeader", "individualHeader", "X-OverridableHeader", "overrideHeader", "X-Envelope", "envelopeHeaderValue"));
    Map<String, String> expectedCookies = ImmutableMap.of("cookie1", "cookie1Value", "cookie2", "cookie2Value");
    assertEquals(headers.size(), expectedHeaders.size(), "Incorrect numnber of headers, found:" + headers.toString());
    assertEquals(cookies.size(), expectedCookies.size(), "Incorrect numnber of cookies, found:" + cookies.toString());
    assertEquals(headersSeenInMuxFilter.size(), expectedHeaders.size(), "Incorrect numnber of headers seen by Mux filter, found:" + headers.toString());
    for (Map.Entry<String, String> header : headers.entrySet()) {
        assertEquals(header.getValue(), expectedHeaders.get(header.getKey()), "Incorrect header value for header: " + header.getKey());
    }
    for (Map.Entry<String, String> cookie : cookies.entrySet()) {
        assertEquals(cookie.getValue(), expectedCookies.get(cookie.getKey()), "Incorrect cookie value for cookie: " + cookie.getKey());
    }
    for (Map.Entry<String, String> header : headersSeenInMuxFilter.entrySet()) {
        assertEquals(header.getValue(), expectedHeaders.get(header.getKey()), "Incorrect header value for header seen by Mux Filter: " + header.getKey());
    }
}