Java tutorial
/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You 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 org.apache.gobblin.writer; import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.Queue; import org.apache.http.client.methods.CloseableHttpResponse; import org.apache.http.client.methods.HttpUriRequest; import org.testng.annotations.Test; import com.google.common.base.Joiner; import com.google.common.collect.ImmutableMap; import com.typesafe.config.Config; import com.typesafe.config.ConfigFactory; import junit.framework.Assert; import lombok.extern.slf4j.Slf4j; import org.apache.gobblin.async.AsyncRequest; import org.apache.gobblin.async.AsyncRequestBuilder; import org.apache.gobblin.async.BufferedRecord; import org.apache.gobblin.async.Callback; import org.apache.gobblin.broker.BrokerConstants; import org.apache.gobblin.broker.SharedResourcesBrokerFactory; import org.apache.gobblin.broker.SharedResourcesBrokerImpl; import org.apache.gobblin.broker.SimpleScopeType; import org.apache.gobblin.broker.iface.SharedResourcesBroker; import org.apache.gobblin.configuration.WorkUnitState; import org.apache.gobblin.http.HttpClient; import org.apache.gobblin.http.ResponseHandler; import org.apache.gobblin.http.ResponseStatus; import org.apache.gobblin.http.StatusType; import org.apache.gobblin.http.ThrottledHttpClient; import org.apache.gobblin.net.Request; import org.apache.gobblin.util.limiter.RateBasedLimiter; import org.apache.gobblin.util.limiter.broker.SharedLimiterFactory; @Test @Slf4j public class AsyncHttpWriterTest { /** * Test successful writes of 4 records */ @Test public void testSuccessfulWrites() { MockHttpClient client = new MockHttpClient(); MockRequestBuilder requestBuilder = new MockRequestBuilder(); MockResponseHandler responseHandler = new MockResponseHandler(); MockAsyncHttpWriterBuilder builder = new MockAsyncHttpWriterBuilder(client, requestBuilder, responseHandler); TestAsyncHttpWriter asyncHttpWriter = new TestAsyncHttpWriter(builder); List<MockWriteCallback> callbacks = new ArrayList<>(); for (int i = 0; i < 4; i++) { callbacks.add(new MockWriteCallback()); } asyncHttpWriter.write(new Object(), callbacks.get(0)); asyncHttpWriter.write(new Object(), callbacks.get(1)); asyncHttpWriter.write(new Object(), callbacks.get(2)); try { asyncHttpWriter.flush(); } catch (IOException e) { Assert.fail("Flush failed"); } asyncHttpWriter.write(new Object(), callbacks.get(3)); try { asyncHttpWriter.close(); } catch (IOException e) { Assert.fail("Close failed"); } // Assert all successful callbacks are invoked for (MockWriteCallback callback : callbacks) { Assert.assertTrue(callback.isSuccess); } Assert.assertTrue(client.isCloseCalled); } @Test public void testSuccessfulWritesWithLimiter() { MockThrottledHttpClient client = new MockThrottledHttpClient(createMockBroker()); MockRequestBuilder requestBuilder = new MockRequestBuilder(); MockResponseHandler responseHandler = new MockResponseHandler(); MockAsyncHttpWriterBuilder builder = new MockAsyncHttpWriterBuilder(client, requestBuilder, responseHandler); TestAsyncHttpWriter asyncHttpWriter = new TestAsyncHttpWriter(builder); List<MockWriteCallback> callbacks = new ArrayList<>(); for (int i = 0; i < 50; i++) { MockWriteCallback callback = new MockWriteCallback(); callbacks.add(callback); asyncHttpWriter.write(new Object(), callback); } try { asyncHttpWriter.close(); } catch (IOException e) { Assert.fail("Close failed"); } // Assert all successful callbacks are invoked for (MockWriteCallback callback : callbacks) { Assert.assertTrue(callback.isSuccess); } Assert.assertTrue(client.getSendTimer().getCount() == 50); Assert.assertTrue(client.isCloseCalled); } private static SharedResourcesBroker createMockBroker() { Joiner JOINER = Joiner.on("."); Config config = ConfigFactory.parseMap(ImmutableMap.of( JOINER.join(BrokerConstants.GOBBLIN_BROKER_CONFIG_PREFIX, SharedLimiterFactory.NAME, SharedLimiterFactory.LIMITER_CLASS_KEY), "qps", JOINER.join(BrokerConstants.GOBBLIN_BROKER_CONFIG_PREFIX, SharedLimiterFactory.NAME, RateBasedLimiter.Factory.QPS_KEY), "10")); SharedResourcesBrokerImpl broker = SharedResourcesBrokerFactory .<SimpleScopeType>createDefaultTopLevelBroker(config, SimpleScopeType.GLOBAL.defaultScopeInstance()); return broker; } /** * Test failure triggered by client error. No retries */ public void testClientError() { MockHttpClient client = new MockHttpClient(); MockRequestBuilder requestBuilder = new MockRequestBuilder(); MockResponseHandler responseHandler = new MockResponseHandler(); MockAsyncHttpWriterBuilder builder = new MockAsyncHttpWriterBuilder(client, requestBuilder, responseHandler); TestAsyncHttpWriter asyncHttpWriter = new TestAsyncHttpWriter(builder); responseHandler.type = StatusType.CLIENT_ERROR; MockWriteCallback callback = new MockWriteCallback(); asyncHttpWriter.write(new Object(), callback); boolean hasAnException = false; try { asyncHttpWriter.close(); } catch (Exception e) { hasAnException = true; } Assert.assertTrue(hasAnException); Assert.assertFalse(callback.isSuccess); Assert.assertTrue(client.isCloseCalled); // No retries are done Assert.assertTrue(client.attempts == 1); Assert.assertTrue(responseHandler.attempts == 1); } /** * Test max attempts triggered by failing to send. Attempt 3 times */ public void testMaxAttempts() { MockHttpClient client = new MockHttpClient(); MockRequestBuilder requestBuilder = new MockRequestBuilder(); MockResponseHandler responseHandler = new MockResponseHandler(); MockAsyncHttpWriterBuilder builder = new MockAsyncHttpWriterBuilder(client, requestBuilder, responseHandler); TestAsyncHttpWriter asyncHttpWriter = new TestAsyncHttpWriter(builder); client.shouldSendSucceed = false; MockWriteCallback callback = new MockWriteCallback(); asyncHttpWriter.write(new Object(), callback); boolean hasAnException = false; try { asyncHttpWriter.close(); } catch (Exception e) { hasAnException = true; } Assert.assertTrue(hasAnException); Assert.assertFalse(callback.isSuccess); Assert.assertTrue(client.isCloseCalled); Assert.assertTrue(client.attempts == AsyncHttpWriter.DEFAULT_MAX_ATTEMPTS); Assert.assertTrue(responseHandler.attempts == 0); } /** * Test server error. Attempt 3 times */ public void testServerError() { MockHttpClient client = new MockHttpClient(); MockRequestBuilder requestBuilder = new MockRequestBuilder(); MockResponseHandler responseHandler = new MockResponseHandler(); MockAsyncHttpWriterBuilder builder = new MockAsyncHttpWriterBuilder(client, requestBuilder, responseHandler); TestAsyncHttpWriter asyncHttpWriter = new TestAsyncHttpWriter(builder); responseHandler.type = StatusType.SERVER_ERROR; MockWriteCallback callback = new MockWriteCallback(); asyncHttpWriter.write(new Object(), callback); boolean hasAnException = false; try { asyncHttpWriter.close(); } catch (Exception e) { hasAnException = true; } Assert.assertTrue(hasAnException); Assert.assertFalse(callback.isSuccess); Assert.assertTrue(client.isCloseCalled); Assert.assertTrue(client.attempts == AsyncHttpWriter.DEFAULT_MAX_ATTEMPTS); Assert.assertTrue(responseHandler.attempts == AsyncHttpWriter.DEFAULT_MAX_ATTEMPTS); } class MockHttpClient implements HttpClient<HttpUriRequest, CloseableHttpResponse> { boolean isCloseCalled = false; int attempts = 0; boolean shouldSendSucceed = true; @Override public CloseableHttpResponse sendRequest(HttpUriRequest request) throws IOException { attempts++; if (shouldSendSucceed) { // We won't consume the response anyway return null; } throw new IOException("Send failed"); } @Override public void sendAsyncRequest(HttpUriRequest request, Callback<CloseableHttpResponse> callback) { throw new UnsupportedOperationException(); } @Override public void close() throws IOException { isCloseCalled = true; } } class MockThrottledHttpClient extends ThrottledHttpClient<HttpUriRequest, CloseableHttpResponse> { boolean isCloseCalled = false; int attempts = 0; boolean shouldSendSucceed = true; public MockThrottledHttpClient(SharedResourcesBroker broker) { super(broker, "resource"); } @Override public CloseableHttpResponse sendRequestImpl(HttpUriRequest request) throws IOException { attempts++; if (shouldSendSucceed) { // We won't consume the response anyway return null; } throw new IOException("Send failed"); } @Override public void close() throws IOException { isCloseCalled = true; } @Override public void sendAsyncRequestImpl(HttpUriRequest request, Callback callback) { throw new UnsupportedOperationException(); } } class MockRequestBuilder implements AsyncRequestBuilder<Object, HttpUriRequest> { @Override public AsyncRequest<Object, HttpUriRequest> buildRequest(Queue<BufferedRecord<Object>> buffer) { BufferedRecord<Object> item = buffer.poll(); AsyncRequest<Object, HttpUriRequest> request = new AsyncRequest<>(); request.markRecord(item, 1); request.setRawRequest(null); return request; } } class MockResponseHandler implements ResponseHandler<HttpUriRequest, CloseableHttpResponse> { volatile StatusType type = StatusType.OK; int attempts = 0; @Override public ResponseStatus handleResponse(Request<HttpUriRequest> request, CloseableHttpResponse response) { attempts++; switch (type) { case OK: return new ResponseStatus(StatusType.OK); case CLIENT_ERROR: return new ResponseStatus(StatusType.CLIENT_ERROR); case SERVER_ERROR: return new ResponseStatus(StatusType.SERVER_ERROR); } return null; } } class MockWriteCallback implements WriteCallback<Object> { boolean isSuccess = false; @Override public void onSuccess(WriteResponse<Object> writeResponse) { isSuccess = true; } @Override public void onFailure(Throwable throwable) { isSuccess = false; } } class MockAsyncHttpWriterBuilder extends AsyncHttpWriterBuilder<Object, HttpUriRequest, CloseableHttpResponse> { MockAsyncHttpWriterBuilder(HttpClient client, MockRequestBuilder requestBuilder, MockResponseHandler responseHandler) { this.client = client; this.asyncRequestBuilder = requestBuilder; this.responseHandler = responseHandler; this.state = new WorkUnitState(); this.queueCapacity = 2; } @Override public DataWriter<Object> build() throws IOException { return null; } @Override public AsyncHttpWriterBuilder<Object, HttpUriRequest, CloseableHttpResponse> fromConfig(Config config) { return null; } } class TestAsyncHttpWriter extends AsyncHttpWriter<Object, HttpUriRequest, CloseableHttpResponse> { public TestAsyncHttpWriter(AsyncHttpWriterBuilder builder) { super(builder); } } }