@RunWith(PowerMockRunner.class) @PrepareForTest(KafkaOffsetBackingStore.class) @PowerMockIgnore("javax.management.*") public class KafkaOffsetBackingStoreTest { private static final String TOPIC = "connect-offsets"; private static final Map<String, String> DEFAULT_PROPS = new HashMap<>(); static { DEFAULT_PROPS.put(KafkaOffsetBackingStore.OFFSET_STORAGE_TOPIC_CONFIG, TOPIC); DEFAULT_PROPS.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, "broker1:9092,broker2:9093"); } private static final Map<ByteBuffer, ByteBuffer> FIRST_SET = new HashMap<>(); static { FIRST_SET.put(buffer("key"), buffer("value")); FIRST_SET.put(null, null); } private static final ByteBuffer TP0_KEY = buffer("TP0KEY"); private static final ByteBuffer TP1_KEY = buffer("TP1KEY"); private static final ByteBuffer TP2_KEY = buffer("TP2KEY"); private static final ByteBuffer TP0_VALUE = buffer("VAL0"); private static final ByteBuffer TP1_VALUE = buffer("VAL1"); private static final ByteBuffer TP2_VALUE = buffer("VAL2"); private static final ByteBuffer TP0_VALUE_NEW = buffer("VAL0_NEW"); private static final ByteBuffer TP1_VALUE_NEW = buffer("VAL1_NEW"); @Mock KafkaBasedLog<byte[], byte[]> storeLog; private KafkaOffsetBackingStore store; private Capture<String> capturedTopic = EasyMock.newCapture(); private Capture<Map<String, Object>> capturedProducerProps = EasyMock.newCapture(); private Capture<Map<String, Object>> capturedConsumerProps = EasyMock.newCapture(); private Capture<Callback<ConsumerRecord<byte[], byte[]>>> capturedConsumedCallback = EasyMock.newCapture(); @Before public void setUp() throws Exception { store = PowerMock.createPartialMockAndInvokeDefaultConstructor( KafkaOffsetBackingStore.class, new String[] {"createKafkaBasedLog"}); } @Test(expected = ConnectException.class) public void testMissingTopic() { store = new KafkaOffsetBackingStore(); store.configure(Collections.<String, Object>emptyMap()); } @Test public void testStartStop() throws Exception { expectConfigure(); expectStart(Collections.EMPTY_LIST); expectStop(); PowerMock.replayAll(); store.configure(DEFAULT_PROPS); assertEquals(TOPIC, capturedTopic.getValue()); assertEquals( "org.apache.kafka.common.serialization.ByteArraySerializer", capturedProducerProps.getValue().get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG)); assertEquals( "org.apache.kafka.common.serialization.ByteArraySerializer", capturedProducerProps.getValue().get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG)); assertEquals( "org.apache.kafka.common.serialization.ByteArrayDeserializer", capturedConsumerProps.getValue().get(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG)); assertEquals( "org.apache.kafka.common.serialization.ByteArrayDeserializer", capturedConsumerProps.getValue().get(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG)); store.start(); store.stop(); PowerMock.verifyAll(); } @Test public void testReloadOnStart() throws Exception { expectConfigure(); expectStart( Arrays.asList( new ConsumerRecord<>(TOPIC, 0, 0, TP0_KEY.array(), TP0_VALUE.array()), new ConsumerRecord<>(TOPIC, 1, 0, TP1_KEY.array(), TP1_VALUE.array()), new ConsumerRecord<>(TOPIC, 0, 1, TP0_KEY.array(), TP0_VALUE_NEW.array()), new ConsumerRecord<>(TOPIC, 1, 1, TP1_KEY.array(), TP1_VALUE_NEW.array()))); expectStop(); PowerMock.replayAll(); store.configure(DEFAULT_PROPS); store.start(); HashMap<ByteBuffer, ByteBuffer> data = Whitebox.getInternalState(store, "data"); assertEquals(TP0_VALUE_NEW, data.get(TP0_KEY)); assertEquals(TP1_VALUE_NEW, data.get(TP1_KEY)); store.stop(); PowerMock.verifyAll(); } @Test public void testGetSet() throws Exception { expectConfigure(); expectStart(Collections.EMPTY_LIST); expectStop(); // First get() against an empty store final Capture<Callback<Void>> firstGetReadToEndCallback = EasyMock.newCapture(); storeLog.readToEnd(EasyMock.capture(firstGetReadToEndCallback)); PowerMock.expectLastCall() .andAnswer( new IAnswer<Object>() { @Override public Object answer() throws Throwable { firstGetReadToEndCallback.getValue().onCompletion(null, null); return null; } }); // Set offsets Capture<org.apache.kafka.clients.producer.Callback> callback0 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP0_KEY.array()), EasyMock.aryEq(TP0_VALUE.array()), EasyMock.capture(callback0)); PowerMock.expectLastCall(); Capture<org.apache.kafka.clients.producer.Callback> callback1 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP1_KEY.array()), EasyMock.aryEq(TP1_VALUE.array()), EasyMock.capture(callback1)); PowerMock.expectLastCall(); // Second get() should get the produced data and return the new values final Capture<Callback<Void>> secondGetReadToEndCallback = EasyMock.newCapture(); storeLog.readToEnd(EasyMock.capture(secondGetReadToEndCallback)); PowerMock.expectLastCall() .andAnswer( new IAnswer<Object>() { @Override public Object answer() throws Throwable { capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 0, 0, TP0_KEY.array(), TP0_VALUE.array())); capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 1, 0, TP1_KEY.array(), TP1_VALUE.array())); secondGetReadToEndCallback.getValue().onCompletion(null, null); return null; } }); // Third get() should pick up data produced by someone else and return those values final Capture<Callback<Void>> thirdGetReadToEndCallback = EasyMock.newCapture(); storeLog.readToEnd(EasyMock.capture(thirdGetReadToEndCallback)); PowerMock.expectLastCall() .andAnswer( new IAnswer<Object>() { @Override public Object answer() throws Throwable { capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 0, 1, TP0_KEY.array(), TP0_VALUE_NEW.array())); capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 1, 1, TP1_KEY.array(), TP1_VALUE_NEW.array())); thirdGetReadToEndCallback.getValue().onCompletion(null, null); return null; } }); PowerMock.replayAll(); store.configure(DEFAULT_PROPS); store.start(); // Getting from empty store should return nulls final AtomicBoolean getInvokedAndPassed = new AtomicBoolean(false); store .get( Arrays.asList(TP0_KEY, TP1_KEY), new Callback<Map<ByteBuffer, ByteBuffer>>() { @Override public void onCompletion(Throwable error, Map<ByteBuffer, ByteBuffer> result) { // Since we didn't read them yet, these will be null assertEquals(null, result.get(TP0_KEY)); assertEquals(null, result.get(TP1_KEY)); getInvokedAndPassed.set(true); } }) .get(10000, TimeUnit.MILLISECONDS); assertTrue(getInvokedAndPassed.get()); // Set some offsets Map<ByteBuffer, ByteBuffer> toSet = new HashMap<>(); toSet.put(TP0_KEY, TP0_VALUE); toSet.put(TP1_KEY, TP1_VALUE); final AtomicBoolean invoked = new AtomicBoolean(false); Future<Void> setFuture = store.set( toSet, new Callback<Void>() { @Override public void onCompletion(Throwable error, Void result) { invoked.set(true); } }); assertFalse(setFuture.isDone()); // Out of order callbacks shouldn't matter, should still require all to be invoked before // invoking the callback // for the store's set callback callback1.getValue().onCompletion(null, null); assertFalse(invoked.get()); callback0.getValue().onCompletion(null, null); setFuture.get(10000, TimeUnit.MILLISECONDS); assertTrue(invoked.get()); // Getting data should read to end of our published data and return it final AtomicBoolean secondGetInvokedAndPassed = new AtomicBoolean(false); store .get( Arrays.asList(TP0_KEY, TP1_KEY), new Callback<Map<ByteBuffer, ByteBuffer>>() { @Override public void onCompletion(Throwable error, Map<ByteBuffer, ByteBuffer> result) { assertEquals(TP0_VALUE, result.get(TP0_KEY)); assertEquals(TP1_VALUE, result.get(TP1_KEY)); secondGetInvokedAndPassed.set(true); } }) .get(10000, TimeUnit.MILLISECONDS); assertTrue(secondGetInvokedAndPassed.get()); // Getting data should read to end of our published data and return it final AtomicBoolean thirdGetInvokedAndPassed = new AtomicBoolean(false); store .get( Arrays.asList(TP0_KEY, TP1_KEY), new Callback<Map<ByteBuffer, ByteBuffer>>() { @Override public void onCompletion(Throwable error, Map<ByteBuffer, ByteBuffer> result) { assertEquals(TP0_VALUE_NEW, result.get(TP0_KEY)); assertEquals(TP1_VALUE_NEW, result.get(TP1_KEY)); thirdGetInvokedAndPassed.set(true); } }) .get(10000, TimeUnit.MILLISECONDS); assertTrue(thirdGetInvokedAndPassed.get()); store.stop(); PowerMock.verifyAll(); } @Test public void testSetFailure() throws Exception { expectConfigure(); expectStart(Collections.EMPTY_LIST); expectStop(); // Set offsets Capture<org.apache.kafka.clients.producer.Callback> callback0 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP0_KEY.array()), EasyMock.aryEq(TP0_VALUE.array()), EasyMock.capture(callback0)); PowerMock.expectLastCall(); Capture<org.apache.kafka.clients.producer.Callback> callback1 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP1_KEY.array()), EasyMock.aryEq(TP1_VALUE.array()), EasyMock.capture(callback1)); PowerMock.expectLastCall(); Capture<org.apache.kafka.clients.producer.Callback> callback2 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP2_KEY.array()), EasyMock.aryEq(TP2_VALUE.array()), EasyMock.capture(callback2)); PowerMock.expectLastCall(); PowerMock.replayAll(); store.configure(DEFAULT_PROPS); store.start(); // Set some offsets Map<ByteBuffer, ByteBuffer> toSet = new HashMap<>(); toSet.put(TP0_KEY, TP0_VALUE); toSet.put(TP1_KEY, TP1_VALUE); toSet.put(TP2_KEY, TP2_VALUE); final AtomicBoolean invoked = new AtomicBoolean(false); final AtomicBoolean invokedFailure = new AtomicBoolean(false); Future<Void> setFuture = store.set( toSet, new Callback<Void>() { @Override public void onCompletion(Throwable error, Void result) { invoked.set(true); if (error != null) invokedFailure.set(true); } }); assertFalse(setFuture.isDone()); // Out of order callbacks shouldn't matter, should still require all to be invoked before // invoking the callback // for the store's set callback callback1.getValue().onCompletion(null, null); assertFalse(invoked.get()); callback2.getValue().onCompletion(null, new KafkaException("bogus error")); assertTrue(invoked.get()); assertTrue(invokedFailure.get()); callback0.getValue().onCompletion(null, null); try { setFuture.get(10000, TimeUnit.MILLISECONDS); fail( "Should have seen KafkaException thrown when waiting on KafkaOffsetBackingStore.set() future"); } catch (ExecutionException e) { // expected assertNotNull(e.getCause()); assertTrue(e.getCause() instanceof KafkaException); } store.stop(); PowerMock.verifyAll(); } private void expectConfigure() throws Exception { PowerMock.expectPrivate( store, "createKafkaBasedLog", EasyMock.capture(capturedTopic), EasyMock.capture(capturedProducerProps), EasyMock.capture(capturedConsumerProps), EasyMock.capture(capturedConsumedCallback)) .andReturn(storeLog); } private void expectStart(final List<ConsumerRecord<byte[], byte[]>> preexistingRecords) throws Exception { storeLog.start(); PowerMock.expectLastCall() .andAnswer( new IAnswer<Object>() { @Override public Object answer() throws Throwable { for (ConsumerRecord<byte[], byte[]> rec : preexistingRecords) capturedConsumedCallback.getValue().onCompletion(null, rec); return null; } }); } private void expectStop() { storeLog.stop(); PowerMock.expectLastCall(); } private static ByteBuffer buffer(String v) { return ByteBuffer.wrap(v.getBytes()); } }
@Test public void testSetFailure() throws Exception { expectConfigure(); expectStart(Collections.EMPTY_LIST); expectStop(); // Set offsets Capture<org.apache.kafka.clients.producer.Callback> callback0 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP0_KEY.array()), EasyMock.aryEq(TP0_VALUE.array()), EasyMock.capture(callback0)); PowerMock.expectLastCall(); Capture<org.apache.kafka.clients.producer.Callback> callback1 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP1_KEY.array()), EasyMock.aryEq(TP1_VALUE.array()), EasyMock.capture(callback1)); PowerMock.expectLastCall(); Capture<org.apache.kafka.clients.producer.Callback> callback2 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP2_KEY.array()), EasyMock.aryEq(TP2_VALUE.array()), EasyMock.capture(callback2)); PowerMock.expectLastCall(); PowerMock.replayAll(); store.configure(DEFAULT_PROPS); store.start(); // Set some offsets Map<ByteBuffer, ByteBuffer> toSet = new HashMap<>(); toSet.put(TP0_KEY, TP0_VALUE); toSet.put(TP1_KEY, TP1_VALUE); toSet.put(TP2_KEY, TP2_VALUE); final AtomicBoolean invoked = new AtomicBoolean(false); final AtomicBoolean invokedFailure = new AtomicBoolean(false); Future<Void> setFuture = store.set( toSet, new Callback<Void>() { @Override public void onCompletion(Throwable error, Void result) { invoked.set(true); if (error != null) invokedFailure.set(true); } }); assertFalse(setFuture.isDone()); // Out of order callbacks shouldn't matter, should still require all to be invoked before // invoking the callback // for the store's set callback callback1.getValue().onCompletion(null, null); assertFalse(invoked.get()); callback2.getValue().onCompletion(null, new KafkaException("bogus error")); assertTrue(invoked.get()); assertTrue(invokedFailure.get()); callback0.getValue().onCompletion(null, null); try { setFuture.get(10000, TimeUnit.MILLISECONDS); fail( "Should have seen KafkaException thrown when waiting on KafkaOffsetBackingStore.set() future"); } catch (ExecutionException e) { // expected assertNotNull(e.getCause()); assertTrue(e.getCause() instanceof KafkaException); } store.stop(); PowerMock.verifyAll(); }
@Test public void testGetSet() throws Exception { expectConfigure(); expectStart(Collections.EMPTY_LIST); expectStop(); // First get() against an empty store final Capture<Callback<Void>> firstGetReadToEndCallback = EasyMock.newCapture(); storeLog.readToEnd(EasyMock.capture(firstGetReadToEndCallback)); PowerMock.expectLastCall() .andAnswer( new IAnswer<Object>() { @Override public Object answer() throws Throwable { firstGetReadToEndCallback.getValue().onCompletion(null, null); return null; } }); // Set offsets Capture<org.apache.kafka.clients.producer.Callback> callback0 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP0_KEY.array()), EasyMock.aryEq(TP0_VALUE.array()), EasyMock.capture(callback0)); PowerMock.expectLastCall(); Capture<org.apache.kafka.clients.producer.Callback> callback1 = EasyMock.newCapture(); storeLog.send( EasyMock.aryEq(TP1_KEY.array()), EasyMock.aryEq(TP1_VALUE.array()), EasyMock.capture(callback1)); PowerMock.expectLastCall(); // Second get() should get the produced data and return the new values final Capture<Callback<Void>> secondGetReadToEndCallback = EasyMock.newCapture(); storeLog.readToEnd(EasyMock.capture(secondGetReadToEndCallback)); PowerMock.expectLastCall() .andAnswer( new IAnswer<Object>() { @Override public Object answer() throws Throwable { capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 0, 0, TP0_KEY.array(), TP0_VALUE.array())); capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 1, 0, TP1_KEY.array(), TP1_VALUE.array())); secondGetReadToEndCallback.getValue().onCompletion(null, null); return null; } }); // Third get() should pick up data produced by someone else and return those values final Capture<Callback<Void>> thirdGetReadToEndCallback = EasyMock.newCapture(); storeLog.readToEnd(EasyMock.capture(thirdGetReadToEndCallback)); PowerMock.expectLastCall() .andAnswer( new IAnswer<Object>() { @Override public Object answer() throws Throwable { capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 0, 1, TP0_KEY.array(), TP0_VALUE_NEW.array())); capturedConsumedCallback .getValue() .onCompletion( null, new ConsumerRecord<>(TOPIC, 1, 1, TP1_KEY.array(), TP1_VALUE_NEW.array())); thirdGetReadToEndCallback.getValue().onCompletion(null, null); return null; } }); PowerMock.replayAll(); store.configure(DEFAULT_PROPS); store.start(); // Getting from empty store should return nulls final AtomicBoolean getInvokedAndPassed = new AtomicBoolean(false); store .get( Arrays.asList(TP0_KEY, TP1_KEY), new Callback<Map<ByteBuffer, ByteBuffer>>() { @Override public void onCompletion(Throwable error, Map<ByteBuffer, ByteBuffer> result) { // Since we didn't read them yet, these will be null assertEquals(null, result.get(TP0_KEY)); assertEquals(null, result.get(TP1_KEY)); getInvokedAndPassed.set(true); } }) .get(10000, TimeUnit.MILLISECONDS); assertTrue(getInvokedAndPassed.get()); // Set some offsets Map<ByteBuffer, ByteBuffer> toSet = new HashMap<>(); toSet.put(TP0_KEY, TP0_VALUE); toSet.put(TP1_KEY, TP1_VALUE); final AtomicBoolean invoked = new AtomicBoolean(false); Future<Void> setFuture = store.set( toSet, new Callback<Void>() { @Override public void onCompletion(Throwable error, Void result) { invoked.set(true); } }); assertFalse(setFuture.isDone()); // Out of order callbacks shouldn't matter, should still require all to be invoked before // invoking the callback // for the store's set callback callback1.getValue().onCompletion(null, null); assertFalse(invoked.get()); callback0.getValue().onCompletion(null, null); setFuture.get(10000, TimeUnit.MILLISECONDS); assertTrue(invoked.get()); // Getting data should read to end of our published data and return it final AtomicBoolean secondGetInvokedAndPassed = new AtomicBoolean(false); store .get( Arrays.asList(TP0_KEY, TP1_KEY), new Callback<Map<ByteBuffer, ByteBuffer>>() { @Override public void onCompletion(Throwable error, Map<ByteBuffer, ByteBuffer> result) { assertEquals(TP0_VALUE, result.get(TP0_KEY)); assertEquals(TP1_VALUE, result.get(TP1_KEY)); secondGetInvokedAndPassed.set(true); } }) .get(10000, TimeUnit.MILLISECONDS); assertTrue(secondGetInvokedAndPassed.get()); // Getting data should read to end of our published data and return it final AtomicBoolean thirdGetInvokedAndPassed = new AtomicBoolean(false); store .get( Arrays.asList(TP0_KEY, TP1_KEY), new Callback<Map<ByteBuffer, ByteBuffer>>() { @Override public void onCompletion(Throwable error, Map<ByteBuffer, ByteBuffer> result) { assertEquals(TP0_VALUE_NEW, result.get(TP0_KEY)); assertEquals(TP1_VALUE_NEW, result.get(TP1_KEY)); thirdGetInvokedAndPassed.set(true); } }) .get(10000, TimeUnit.MILLISECONDS); assertTrue(thirdGetInvokedAndPassed.get()); store.stop(); PowerMock.verifyAll(); }
@Test public void testRun() throws Exception { HttpClient httpClient = EasyMock.createMock(HttpClient.class); final URL url = new URL("http://foo/druid/v2/"); SettableFuture<InputStream> futureResult = SettableFuture.create(); Capture<Request> capturedRequest = EasyMock.newCapture(); EasyMock.expect( httpClient.go( EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject())) .andReturn(futureResult) .times(1); SettableFuture futureException = SettableFuture.create(); EasyMock.expect( httpClient.go( EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject())) .andReturn(futureException) .times(1); EasyMock.expect( httpClient.go( EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject())) .andReturn(SettableFuture.create()) .atLeastOnce(); EasyMock.replay(httpClient); final ServerSelector serverSelector = new ServerSelector( new DataSegment( "test", new Interval("2013-01-01/2013-01-02"), new DateTime("2013-01-01").toString(), Maps.<String, Object>newHashMap(), Lists.<String>newArrayList(), Lists.<String>newArrayList(), NoneShardSpec.instance(), 0, 0L), new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy())); DirectDruidClient client1 = new DirectDruidClient( new ReflectionQueryToolChestWarehouse(), QueryRunnerTestHelper.NOOP_QUERYWATCHER, new DefaultObjectMapper(), httpClient, "foo", new NoopServiceEmitter()); DirectDruidClient client2 = new DirectDruidClient( new ReflectionQueryToolChestWarehouse(), QueryRunnerTestHelper.NOOP_QUERYWATCHER, new DefaultObjectMapper(), httpClient, "foo2", new NoopServiceEmitter()); QueryableDruidServer queryableDruidServer1 = new QueryableDruidServer( new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0), client1); serverSelector.addServerAndUpdateSegment(queryableDruidServer1, serverSelector.getSegment()); QueryableDruidServer queryableDruidServer2 = new QueryableDruidServer( new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0), client2); serverSelector.addServerAndUpdateSegment(queryableDruidServer2, serverSelector.getSegment()); TimeBoundaryQuery query = Druids.newTimeBoundaryQueryBuilder().dataSource("test").build(); HashMap<String, List> context = Maps.newHashMap(); Sequence s1 = client1.run(query, context); Assert.assertTrue(capturedRequest.hasCaptured()); Assert.assertEquals(url, capturedRequest.getValue().getUrl()); Assert.assertEquals(HttpMethod.POST, capturedRequest.getValue().getMethod()); Assert.assertEquals(1, client1.getNumOpenConnections()); // simulate read timeout Sequence s2 = client1.run(query, context); Assert.assertEquals(2, client1.getNumOpenConnections()); futureException.setException(new ReadTimeoutException()); Assert.assertEquals(1, client1.getNumOpenConnections()); // subsequent connections should work Sequence s3 = client1.run(query, context); Sequence s4 = client1.run(query, context); Sequence s5 = client1.run(query, context); Assert.assertTrue(client1.getNumOpenConnections() == 4); // produce result for first connection futureResult.set( new ByteArrayInputStream( "[{\"timestamp\":\"2014-01-01T01:02:03Z\", \"result\": 42.0}]".getBytes())); List<Result> results = Sequences.toList(s1, Lists.<Result>newArrayList()); Assert.assertEquals(1, results.size()); Assert.assertEquals(new DateTime("2014-01-01T01:02:03Z"), results.get(0).getTimestamp()); Assert.assertEquals(3, client1.getNumOpenConnections()); client2.run(query, context); client2.run(query, context); Assert.assertTrue(client2.getNumOpenConnections() == 2); Assert.assertTrue(serverSelector.pick() == queryableDruidServer2); EasyMock.verify(httpClient); }
@Test public void testQueryInterruptionExceptionLogMessage() throws JsonProcessingException { HttpClient httpClient = EasyMock.createMock(HttpClient.class); SettableFuture<Object> interruptionFuture = SettableFuture.create(); Capture<Request> capturedRequest = EasyMock.newCapture(); String hostName = "localhost:8080"; EasyMock.expect( httpClient.go( EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject())) .andReturn(interruptionFuture) .anyTimes(); EasyMock.replay(httpClient); DataSegment dataSegment = new DataSegment( "test", new Interval("2013-01-01/2013-01-02"), new DateTime("2013-01-01").toString(), Maps.<String, Object>newHashMap(), Lists.<String>newArrayList(), Lists.<String>newArrayList(), NoneShardSpec.instance(), 0, 0L); final ServerSelector serverSelector = new ServerSelector( dataSegment, new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy())); DirectDruidClient client1 = new DirectDruidClient( new ReflectionQueryToolChestWarehouse(), QueryRunnerTestHelper.NOOP_QUERYWATCHER, new DefaultObjectMapper(), httpClient, hostName, new NoopServiceEmitter()); QueryableDruidServer queryableDruidServer = new QueryableDruidServer( new DruidServer("test1", hostName, 0, "historical", DruidServer.DEFAULT_TIER, 0), client1); serverSelector.addServerAndUpdateSegment(queryableDruidServer, dataSegment); TimeBoundaryQuery query = Druids.newTimeBoundaryQueryBuilder().dataSource("test").build(); HashMap<String, List> context = Maps.newHashMap(); interruptionFuture.set(new ByteArrayInputStream("{\"error\":\"testing\"}".getBytes())); Sequence results = client1.run(query, context); QueryInterruptedException actualException = null; try { Sequences.toList(results, Lists.newArrayList()); } catch (QueryInterruptedException e) { actualException = e; } Assert.assertNotNull(actualException); Assert.assertEquals(actualException.getMessage(), QueryInterruptedException.UNKNOWN_EXCEPTION); Assert.assertEquals(actualException.getCauseMessage(), "testing"); Assert.assertEquals(actualException.getHost(), hostName); EasyMock.verify(httpClient); }
@Test public void testCancel() throws Exception { HttpClient httpClient = EasyMock.createStrictMock(HttpClient.class); Capture<Request> capturedRequest = EasyMock.newCapture(); ListenableFuture<Object> cancelledFuture = Futures.immediateCancelledFuture(); SettableFuture<Object> cancellationFuture = SettableFuture.create(); EasyMock.expect( httpClient.go( EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject())) .andReturn(cancelledFuture) .once(); EasyMock.expect( httpClient.go( EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject())) .andReturn(cancellationFuture) .once(); EasyMock.replay(httpClient); final ServerSelector serverSelector = new ServerSelector( new DataSegment( "test", new Interval("2013-01-01/2013-01-02"), new DateTime("2013-01-01").toString(), Maps.<String, Object>newHashMap(), Lists.<String>newArrayList(), Lists.<String>newArrayList(), NoneShardSpec.instance(), 0, 0L), new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy())); DirectDruidClient client1 = new DirectDruidClient( new ReflectionQueryToolChestWarehouse(), QueryRunnerTestHelper.NOOP_QUERYWATCHER, new DefaultObjectMapper(), httpClient, "foo", new NoopServiceEmitter()); QueryableDruidServer queryableDruidServer1 = new QueryableDruidServer( new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0), client1); serverSelector.addServerAndUpdateSegment(queryableDruidServer1, serverSelector.getSegment()); TimeBoundaryQuery query = Druids.newTimeBoundaryQueryBuilder().dataSource("test").build(); HashMap<String, List> context = Maps.newHashMap(); cancellationFuture.set( new StatusResponseHolder(HttpResponseStatus.OK, new StringBuilder("cancelled"))); Sequence results = client1.run(query, context); Assert.assertEquals(HttpMethod.DELETE, capturedRequest.getValue().getMethod()); Assert.assertEquals(0, client1.getNumOpenConnections()); QueryInterruptedException exception = null; try { Sequences.toList(results, Lists.newArrayList()); } catch (QueryInterruptedException e) { exception = e; } Assert.assertNotNull(exception); EasyMock.verify(httpClient); }