Repository: apache/kafka
File: connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java (Java)

Code before:
        expectStop();

        // First get() against an empty store
        final Capture<Callback<Void>> firstGetReadToEndCallback = EasyMock.newCapture();
        storeLog.readToEnd(EasyMock.capture(firstGetReadToEndCallback));
        PowerMock.expectLastCall().andAnswer(() -> {
            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(() -> {
            capturedConsumedCallback.getValue().onCompletion(null,
                new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0, 0, TP0_KEY.array(), TP0_VALUE.array(),
                    new RecordHeaders(), Optional.empty()));
            capturedConsumedCallback.getValue().onCompletion(null,
                new ConsumerRecord<>(TOPIC, 1, 0, 0L, TimestampType.CREATE_TIME, 0, 0, TP1_KEY.array(), TP1_VALUE.array(),
                    new RecordHeaders(), Optional.empty()));
            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(() -> {
            capturedConsumedCallback.getValue().onCompletion(null,
                new ConsumerRecord<>(TOPIC, 0, 1, 0L, TimestampType.CREATE_TIME, 0, 0, TP0_KEY.array(), TP0_VALUE_NEW.array(),
                    new RecordHeaders(), Optional.empty()));
            capturedConsumedCallback.getValue().onCompletion(null,
                new ConsumerRecord<>(TOPIC, 1, 1, 0L, TimestampType.CREATE_TIME, 0, 0, TP1_KEY.array(), TP1_VALUE_NEW.array(),
                    new RecordHeaders(), Optional.empty()));
            thirdGetReadToEndCallback.getValue().onCompletion(null, null);
            return null;
        });

        expectClusterId();
        PowerMock.replayAll();

        store.configure(DEFAULT_DISTRIBUTED_CONFIG);
        store.start();

Review comment (suggestion):
We don't have to set up the expectations for all three calls to `KafkaOffsetBackingStore::get` at once. We can set up the expectation for each right before we invoke the method.

This has a few advantages:
- Cleanliness (the chained `doAnswer` method invocations are a little unintuitive)
- Readability (placing expectation-setting code closer to execution code makes it easier to see how the two are related)

In general, I think we should try to follow Mockito idioms more closely with these tests by establishing expectations as late as possible in tests, instead of having a monolithic expectation-setting section at the beginning of each test case.

Code after:
        }).when(storeLog).start();

        store.configure(DEFAULT_DISTRIBUTED_CONFIG);
        store.start();

        HashMap<ByteBuffer, ByteBuffer> data = store.data;
        assertEquals(TP0_VALUE_NEW, data.get(TP0_KEY));
        assertEquals(TP1_VALUE_NEW, data.get(TP1_KEY));

        store.stop();

        verify(storeLog).stop();
    }

    @Test
    public void testGetSet() throws Exception {
        store.configure(DEFAULT_DISTRIBUTED_CONFIG);
        store.start();

        verify(storeLog).start();

        doAnswer(invocation -> {
            // First get() against an empty store
            storeLogCallbackArgumentCaptor.getValue().onCompletion(null, null);
            return null;
        }).when(storeLog).readToEnd(storeLogCallbackArgumentCaptor.capture());

        // Getting from empty store should return nulls
        Map<ByteBuffer, ByteBuffer> offsets = store.get(Arrays.asList(TP0_KEY, TP1_KEY)).get(10000, TimeUnit.MILLISECONDS);
        // Since we didn't read them yet, these will be null
        assertNull(offsets.get(TP0_KEY));
        assertNull(offsets.get(TP1_KEY));
        // 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, (error, result) -> invoked.set(true));
        assertFalse(setFuture.isDone());
        // Set offsets
        ArgumentCaptor<org.apache.kafka.clients.producer.Callback> callback0 = ArgumentCaptor.forClass(org.apache.kafka.clients.producer.Callback.class);
        verify(storeLog).send(aryEq(TP0_KEY.array()), aryEq(TP0_VALUE.array()), callback0.capture());
        ArgumentCaptor<org.apache.kafka.clients.producer.Callback> callback1 = ArgumentCaptor.forClass(org.apache.kafka.clients.producer.Callback.class);
        verify(storeLog).send(aryEq(TP1_KEY.array()), aryEq(TP1_VALUE.array()), callback1.capture());
        // 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());