GafferAdder implements Serializable { public void add(final Element element) { if (null == element) { return; } if (null == queue) { queue = new ConsumableBlockingQueue<>(maxQueueSize); restart = true; } try { queue.put(element); } catch (final InterruptedException e) { throw new RuntimeException("Interrupted while waiting to add an element to the queue", e); } if (restart && !queue.isEmpty()) { restart = false; store.runAsync(() -> { try { store.execute(new AddElements.Builder() .input(queue) .validate(validate) .skipInvalidElements(skipInvalid) .build(), new Context(new User())); restart = true; } catch (final OperationException e) { throw new RuntimeException(e.getMessage(), e); } }); } } <OP extends Validatable & Operation> GafferAdder(final OP operation, final Store store); void initialise(); void add(final Element element); static final int MAX_QUEUE_SIZE_DEFAULT; }
@Test public void shouldAddElementsToStore() throws Exception { final AddElementsFromSocket op = mock(AddElementsFromSocket.class); final Store store = mock(Store.class); given(store.getProperties()).willReturn(new StoreProperties()); given(store.getSchema()).willReturn(new Schema()); given(op.isValidate()).willReturn(true); given(op.isSkipInvalidElements()).willReturn(false); given(op.getOption(FlinkConstants.MAX_QUEUE_SIZE)).willReturn(MAX_QUEUE_SIZE_OPTION); final Element element = mock(Element.class); final GafferAdder adder = new GafferAdder(op, store); adder.add(element); final ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class); verify(store).runAsync(runnableCaptor.capture()); runnableCaptor.getValue().run(); ArgumentCaptor<AddElements> opCaptor = ArgumentCaptor.forClass(AddElements.class); verify(store).execute( opCaptor.capture(), Mockito.any() ); final ConsumableBlockingQueue<Element> expectedQueue = new ConsumableBlockingQueue<>(MAX_QUEUE_SIZE_VALUE); expectedQueue.put(element); verify(store).execute(Mockito.eq(new AddElements.Builder() .input(expectedQueue) .validate(true) .skipInvalidElements(false) .build()), Mockito.any()); }
@Test public void shouldRestartAddElementsIfPauseInIngest() throws Exception { final AddElementsFromSocket op = mock(AddElementsFromSocket.class); final Store store = mock(Store.class); given(store.getProperties()).willReturn(new StoreProperties()); given(store.getSchema()).willReturn(new Schema()); given(op.isValidate()).willReturn(true); given(op.isSkipInvalidElements()).willReturn(false); given(op.getOption(FlinkConstants.MAX_QUEUE_SIZE)).willReturn(MAX_QUEUE_SIZE_OPTION); final Element element = mock(Element.class); final Element element2 = mock(Element.class); final GafferAdder adder = new GafferAdder(op, store); adder.add(element); final ArgumentCaptor<Runnable> runnableCaptor1 = ArgumentCaptor.forClass(Runnable.class); verify(store).runAsync(runnableCaptor1.capture()); runnableCaptor1.getValue().run(); final ConsumableBlockingQueue<Element> expectedQueue = new ConsumableBlockingQueue<>(MAX_QUEUE_SIZE_VALUE); expectedQueue.put(element); verify(store).execute(Mockito.eq(new AddElements.Builder() .input(expectedQueue) .validate(true) .skipInvalidElements(false) .build()), Mockito.any()); Mockito.reset(store); adder.add(element2); final ArgumentCaptor<Runnable> runnableCaptor2 = ArgumentCaptor.forClass(Runnable.class); verify(store).runAsync(runnableCaptor2.capture()); runnableCaptor2.getValue().run(); expectedQueue.put(element2); verify(store).execute(Mockito.eq(new AddElements.Builder() .input(expectedQueue) .validate(true) .skipInvalidElements(false) .build()), Mockito.any()); }
@Test public void shouldAddElementsIfInvokeCalledMultipleTimes() throws Exception { final int duplicates = 4; final AddElementsFromSocket op = mock(AddElementsFromSocket.class); final Store store = mock(Store.class); given(store.getProperties()).willReturn(new StoreProperties()); given(store.getSchema()).willReturn(new Schema()); given(op.isValidate()).willReturn(true); given(op.isSkipInvalidElements()).willReturn(false); given(op.getOption(FlinkConstants.MAX_QUEUE_SIZE)).willReturn(MAX_QUEUE_SIZE_OPTION); final Element element = mock(Element.class); final GafferAdder adder = new GafferAdder(op, store); for (int i = 0; i < duplicates; i++) { adder.add(element); } final ArgumentCaptor<Runnable> runnableCaptor1 = ArgumentCaptor.forClass(Runnable.class); verify(store).runAsync(runnableCaptor1.capture()); assertEquals(1, runnableCaptor1.getAllValues().size()); runnableCaptor1.getValue().run(); final ConsumableBlockingQueue<Element> expectedQueue = new ConsumableBlockingQueue<>(MAX_QUEUE_SIZE_VALUE); for (int i = 0; i < duplicates; i++) { expectedQueue.put(element); } verify(store).execute(Mockito.eq(new AddElements.Builder() .input(expectedQueue) .validate(true) .skipInvalidElements(false) .build()), Mockito.any()); }