Java EasyMock.newCapture方法代码示例

import org.easymock.EasyMock; //导入方法依赖的package包/类
private Capture<OFMeterMod> prepareForMeterTest() {
    Capture<OFMeterMod> capture = EasyMock.newCapture();

    expect(ofSwitchService.getSwitch(dpid)).andStubReturn(iofSwitch);
    expect(iofSwitch.getOFFactory()).andStubReturn(ofFactory);
    expect(iofSwitch.getSwitchDescription()).andStubReturn(switchDescription);
    expect(switchDescription.getManufacturerDescription()).andStubReturn("");
    expect(iofSwitch.write(capture(capture))).andReturn(true);
    EasyMock.expectLastCall();

    replay(ofSwitchService);
    replay(iofSwitch);
    replay(switchDescription);

    return capture;
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
private void expectApplyTransformationChain(boolean anyTimes) {
    final Capture<SourceRecord> recordCapture = EasyMock.newCapture();
    IExpectationSetters<SourceRecord> convertKeyExpect = EasyMock.expect(transformationChain.apply(EasyMock.capture(recordCapture)));
    if (anyTimes)
        convertKeyExpect.andStubAnswer(new IAnswer<SourceRecord>() {
            @Override
            public SourceRecord answer() {
                return recordCapture.getValue();
            }
        });
    else
        convertKeyExpect.andAnswer(new IAnswer<SourceRecord>() {
            @Override
            public SourceRecord answer() {
                return recordCapture.getValue();
            }
        });
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@Test
public void testSchedule() throws Exception {
    Capture<Runnable> taskWrapper = EasyMock.newCapture();

    ScheduledFuture commitFuture = PowerMock.createMock(ScheduledFuture.class);
    EasyMock.expect(executor.scheduleWithFixedDelay(
            EasyMock.capture(taskWrapper), eq(DEFAULT_OFFSET_COMMIT_INTERVAL_MS),
            eq(DEFAULT_OFFSET_COMMIT_INTERVAL_MS), eq(TimeUnit.MILLISECONDS))
    ).andReturn(commitFuture);

    ConnectorTaskId taskId = PowerMock.createMock(ConnectorTaskId.class);
    WorkerSourceTask task = PowerMock.createMock(WorkerSourceTask.class);

    EasyMock.expect(committers.put(taskId, commitFuture)).andReturn(null);

    PowerMock.replayAll();

    committer.schedule(taskId, task);
    assertTrue(taskWrapper.hasCaptured());
    assertNotNull(taskWrapper.getValue());

    PowerMock.verifyAll();
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@Test
public void testMissingTimestampPropagation() throws Exception {
    expectInitializeTask();
    expectConsumerPoll(1, RecordBatch.NO_TIMESTAMP, TimestampType.CREATE_TIME);
    expectConversionAndTransformation(1);

    Capture<Collection<SinkRecord>> records = EasyMock.newCapture(CaptureType.ALL);

    sinkTask.put(EasyMock.capture(records));

    PowerMock.replayAll();

    workerTask.initialize(TASK_CONFIG);
    workerTask.initializeAndStart();
    workerTask.iteration();

    SinkRecord record = records.getValue().iterator().next();

    // we expect null for missing timestamp, the sentinel value of Record.NO_TIMESTAMP is Kafka's API
    assertEquals(null, record.timestamp());
    assertEquals(TimestampType.CREATE_TIME, record.timestampType());

    PowerMock.verifyAll();
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
private Capture<OFFlowMod> prepareForInstallTest() {
    Capture<OFFlowMod> capture = EasyMock.newCapture();

    expect(ofSwitchService.getSwitch(dpid)).andStubReturn(iofSwitch);
    expect(iofSwitch.getOFFactory()).andStubReturn(ofFactory);
    expect(iofSwitch.getSwitchDescription()).andStubReturn(switchDescription);
    expect(switchDescription.getManufacturerDescription()).andStubReturn("");
    expect(iofSwitch.write(capture(capture))).andReturn(true);
    EasyMock.expectLastCall();

    replay(ofSwitchService);
    replay(iofSwitch);
    replay(switchDescription);

    return capture;
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@Override
public void setup() {
    super.setup();
    Map<String, String> workerProps = new HashMap<>();
    workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
    workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
    workerProps.put("internal.key.converter", "org.apache.kafka.connect.json.JsonConverter");
    workerProps.put("internal.value.converter", "org.apache.kafka.connect.json.JsonConverter");
    workerProps.put("internal.key.converter.schemas.enable", "false");
    workerProps.put("internal.value.converter.schemas.enable", "false");
    workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
    plugins = new Plugins(workerProps);
    config = new StandaloneConfig(workerProps);
    producerCallbacks = EasyMock.newCapture();
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@Test
public void testPollRedelivery() throws Exception {
    expectInitializeTask();
    expectPollInitialAssignment();

    // If a retriable exception is thrown, we should redeliver the same batch, pausing the consumer in the meantime
    expectConsumerPoll(1);
    expectConversionAndTransformation(1);
    Capture<Collection<SinkRecord>> records = EasyMock.newCapture(CaptureType.ALL);
    sinkTask.put(EasyMock.capture(records));
    EasyMock.expectLastCall().andThrow(new RetriableException("retry"));
    // Pause
    HashSet<TopicPartition> partitions = new HashSet<>(asList(TOPIC_PARTITION, TOPIC_PARTITION2));
    EasyMock.expect(consumer.assignment()).andReturn(partitions);
    consumer.pause(partitions);
    PowerMock.expectLastCall();

    // Retry delivery should succeed
    expectConsumerPoll(0);
    sinkTask.put(EasyMock.capture(records));
    EasyMock.expectLastCall();
    // And unpause
    EasyMock.expect(consumer.assignment()).andReturn(partitions);
    consumer.resume(singleton(TOPIC_PARTITION));
    PowerMock.expectLastCall();
    consumer.resume(singleton(TOPIC_PARTITION2));
    PowerMock.expectLastCall();

    PowerMock.replayAll();

    workerTask.initialize(TASK_CONFIG);
    workerTask.initializeAndStart();
    workerTask.iteration();
    workerTask.iteration();
    workerTask.iteration();

    PowerMock.verifyAll();
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@Test
public void testProducerError() throws Exception {
    expectStart();
    TestFuture<RecordMetadata> tp0Future = new TestFuture<>();
    ProducerRecord<String, String> tp0Record = new ProducerRecord<>(TOPIC, TP0_KEY, TP0_VALUE);
    Capture<org.apache.kafka.clients.producer.Callback> callback0 = EasyMock.newCapture();
    EasyMock.expect(producer.send(EasyMock.eq(tp0Record), EasyMock.capture(callback0))).andReturn(tp0Future);

    expectStop();

    PowerMock.replayAll();

    Map<TopicPartition, Long> endOffsets = new HashMap<>();
    endOffsets.put(TP0, 0L);
    endOffsets.put(TP1, 0L);
    consumer.updateEndOffsets(endOffsets);
    store.start();
    assertEquals(CONSUMER_ASSIGNMENT, consumer.assignment());
    assertEquals(0L, consumer.position(TP0));
    assertEquals(0L, consumer.position(TP1));

    final AtomicReference<Throwable> setException = new AtomicReference<>();
    store.send(TP0_KEY, TP0_VALUE, new org.apache.kafka.clients.producer.Callback() {
        @Override
        public void onCompletion(RecordMetadata metadata, Exception exception) {
            assertNull(setException.get()); // Should only be invoked once
            setException.set(exception);
        }
    });
    KafkaException exc = new LeaderNotAvailableException("Error");
    tp0Future.resolve(exc);
    callback0.getValue().onCompletion(null, exc);
    assertNotNull(setException.get());

    store.stop();

    assertFalse(Whitebox.<Thread>getInternalState(store, "thread").isAlive());
    assertTrue(consumer.closed());
    PowerMock.verifyAll();
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
private void expectConversionAndTransformation(final int numMessages) {
    EasyMock.expect(keyConverter.toConnectData(TOPIC, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY)).times(numMessages);
    EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE)).times(numMessages);

    final Capture<SinkRecord> recordCapture = EasyMock.newCapture();
    EasyMock.expect(transformationChain.apply(EasyMock.capture(recordCapture)))
            .andAnswer(new IAnswer<SinkRecord>() {
                @Override
                public SinkRecord answer() {
                    return recordCapture.getValue();
                }
            }).times(numMessages);
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
private Capture<Collection<SinkRecord>> expectPolls(final long pollDelayMs) throws Exception {
    // Stub out all the consumer stream/iterator responses, which we just want to verify occur,
    // but don't care about the exact details here.
    EasyMock.expect(consumer.poll(EasyMock.anyLong())).andStubAnswer(
            new IAnswer<ConsumerRecords<byte[], byte[]>>() {
                @Override
                public ConsumerRecords<byte[], byte[]> answer() throws Throwable {
                    // "Sleep" so time will progress
                    time.sleep(pollDelayMs);
                    ConsumerRecords<byte[], byte[]> records = new ConsumerRecords<>(
                            Collections.singletonMap(
                                    new TopicPartition(TOPIC, PARTITION),
                                    Arrays.asList(
                                            new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturned, TIMESTAMP, TIMESTAMP_TYPE, 0L, 0, 0, RAW_KEY, RAW_VALUE)
                                    )));
                    recordsReturned++;
                    return records;
                }
            });
    EasyMock.expect(keyConverter.toConnectData(TOPIC, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY)).anyTimes();
    EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE)).anyTimes();

    final Capture<SinkRecord> recordCapture = EasyMock.newCapture();
    EasyMock.expect(transformationChain.apply(EasyMock.capture(recordCapture))).andAnswer(new IAnswer<SinkRecord>() {
        @Override
        public SinkRecord answer() {
            return recordCapture.getValue();
        }
    }).anyTimes();

    Capture<Collection<SinkRecord>> capturedRecords = EasyMock.newCapture(CaptureType.ALL);
    sinkTask.put(EasyMock.capture(capturedRecords));
    EasyMock.expectLastCall().anyTimes();
    return capturedRecords;
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
private Capture<OffsetCommitCallback> expectOffsetCommit(final long expectedMessages,
                                                         final RuntimeException error,
                                                         final Exception consumerCommitError,
                                                         final long consumerCommitDelayMs,
                                                         final boolean invokeCallback)
        throws Exception {
    final long finalOffset = FIRST_OFFSET + expectedMessages;

    // All assigned partitions will have offsets committed, but we've only processed messages/updated offsets for one
    final Map<TopicPartition, OffsetAndMetadata> offsetsToCommit = new HashMap<>();
    offsetsToCommit.put(TOPIC_PARTITION, new OffsetAndMetadata(finalOffset));
    offsetsToCommit.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    offsetsToCommit.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
    sinkTask.preCommit(offsetsToCommit);
    IExpectationSetters<Object> expectation = PowerMock.expectLastCall();
    if (error != null) {
        expectation.andThrow(error).once();
        return null;
    } else {
        expectation.andReturn(offsetsToCommit);
    }

    final Capture<OffsetCommitCallback> capturedCallback = EasyMock.newCapture();
    consumer.commitAsync(EasyMock.eq(offsetsToCommit),
            EasyMock.capture(capturedCallback));
    PowerMock.expectLastCall().andAnswer(new IAnswer<Object>() {
        @Override
        public Object answer() throws Throwable {
            time.sleep(consumerCommitDelayMs);
            if (invokeCallback)
                capturedCallback.getValue().onComplete(offsetsToCommit, consumerCommitError);
            return null;
        }
    });
    return capturedCallback;
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@Test
public void taskStatus() {
    ConnectorTaskId taskId = new ConnectorTaskId("connector", 0);
    String workerId = "workerId";

    ConfigBackingStore configStore = strictMock(ConfigBackingStore.class);
    StatusBackingStore statusStore = strictMock(StatusBackingStore.class);

    AbstractHerder herder = partialMockBuilder(AbstractHerder.class)
            .withConstructor(Worker.class, String.class, StatusBackingStore.class, ConfigBackingStore.class)
            .withArgs(worker, workerId, statusStore, configStore)
            .addMockedMethod("generation")
            .createMock();

    EasyMock.expect(herder.generation()).andStubReturn(5);

    final Capture<TaskStatus> statusCapture = EasyMock.newCapture();
    statusStore.putSafe(EasyMock.capture(statusCapture));
    EasyMock.expectLastCall();

    EasyMock.expect(statusStore.get(taskId)).andAnswer(new IAnswer<TaskStatus>() {
        @Override
        public TaskStatus answer() throws Throwable {
            return statusCapture.getValue();
        }
    });

    replayAll();

    herder.onFailure(taskId, new RuntimeException());

    ConnectorStateInfo.TaskState taskState = herder.taskStatus(taskId);
    assertEquals(workerId, taskState.workerId());
    assertEquals("FAILED", taskState.state());
    assertEquals(0, taskState.id());
    assertNotNull(taskState.trace());

    verifyAll();
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
private Capture<ProducerRecord<byte[], byte[]>> expectSendRecord(boolean anyTimes, boolean isRetry, boolean succeed) throws InterruptedException {
    expectConvertKeyValue(anyTimes);
    expectApplyTransformationChain(anyTimes);

    Capture<ProducerRecord<byte[], byte[]>> sent = EasyMock.newCapture();

    // 1. Offset data is passed to the offset storage.
    if (!isRetry) {
        offsetWriter.offset(PARTITION, OFFSET);
        if (anyTimes)
            PowerMock.expectLastCall().anyTimes();
        else
            PowerMock.expectLastCall();
    }

    // 2. Converted data passed to the producer, which will need callbacks invoked for flush to work
    IExpectationSetters<Future<RecordMetadata>> expect = EasyMock.expect(
            producer.send(EasyMock.capture(sent),
                    EasyMock.capture(producerCallbacks)));
    IAnswer<Future<RecordMetadata>> expectResponse = new IAnswer<Future<RecordMetadata>>() {
        @Override
        public Future<RecordMetadata> answer() throws Throwable {
            synchronized (producerCallbacks) {
                for (org.apache.kafka.clients.producer.Callback cb : producerCallbacks.getValues()) {
                    cb.onCompletion(new RecordMetadata(new TopicPartition("foo", 0), 0, 0,
                                                       0L, 0L, 0, 0), null);
                }
                producerCallbacks.reset();
            }
            return sendFuture;
        }
    };
    if (anyTimes)
        expect.andStubAnswer(expectResponse);
    else
        expect.andAnswer(expectResponse);

    // 3. As a result of a successful producer send callback, we'll notify the source task of the record commit
    expectTaskCommitRecord(anyTimes, succeed);

    return sent;
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@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_DISTRIBUTED_CONFIG);
    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();
}
 

import org.easymock.EasyMock; //导入方法依赖的package包/类
@Test
public void testRequestCommit() throws Exception {
    expectInitializeTask();

    expectPollInitialAssignment();

    expectConsumerPoll(1);
    expectConversionAndTransformation(1);
    sinkTask.put(EasyMock.<Collection<SinkRecord>>anyObject());
    EasyMock.expectLastCall();

    final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
    offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
    offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    sinkTask.preCommit(offsets);
    EasyMock.expectLastCall().andReturn(offsets);

    final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
    consumer.commitAsync(EasyMock.eq(offsets), EasyMock.capture(callback));
    EasyMock.expectLastCall().andAnswer(new IAnswer<Void>() {
        @Override
        public Void answer() throws Throwable {
            callback.getValue().onComplete(offsets, null);
            return null;
        }
    });

    expectConsumerPoll(0);
    sinkTask.put(Collections.<SinkRecord>emptyList());
    EasyMock.expectLastCall();

    PowerMock.replayAll();

    workerTask.initialize(TASK_CONFIG);
    workerTask.initializeAndStart();

    workerTask.iteration(); // initial assignment

    workerTask.iteration(); // first record delivered

    sinkTaskContext.getValue().requestCommit();
    assertTrue(sinkTaskContext.getValue().isCommitRequested());
    assertNotEquals(offsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
    workerTask.iteration(); // triggers the commit
    assertFalse(sinkTaskContext.getValue().isCommitRequested()); // should have been cleared
    assertEquals(offsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
    assertEquals(0, workerTask.commitFailures());

    PowerMock.verifyAll();
}