Java 类org.easymock.IExpectationSetters 实例源码

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();
            }
        });
}

@SuppressWarnings("unchecked")
private void expectOffsetFlush(boolean succeed) throws Exception {
    EasyMock.expect(offsetWriter.beginFlush()).andReturn(true);
    Future<Void> flushFuture = PowerMock.createMock(Future.class);
    EasyMock.expect(offsetWriter.doFlush(EasyMock.anyObject(Callback.class))).andReturn(flushFuture);
    // Should throw for failure
    IExpectationSetters<Void> futureGetExpect = EasyMock.expect(
            flushFuture.get(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class)));
    if (succeed) {
        sourceTask.commit();
        EasyMock.expectLastCall();
        futureGetExpect.andReturn(null);
    } else {
        futureGetExpect.andThrow(new TimeoutException());
        offsetWriter.cancelFlush();
        PowerMock.expectLastCall();
    }
}

@SuppressWarnings("unchecked")
private IExpectationSetters<Object> expectOnePoll() {
    // Currently the SinkTask's put() method will not be invoked unless we provide some data, so instead of
    // returning empty data, we return one record. The expectation is that the data will be ignored by the
    // response behavior specified using the return value of this method.
    EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
            new IAnswer<ConsumerRecords<byte[], byte[]>>() {
                @Override
                public ConsumerRecords<byte[], byte[]> answer() throws Throwable {
                    // "Sleep" so time will progress
                    time.sleep(1L);
                    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));
    EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE));
    sinkTask.put(EasyMock.anyObject(Collection.class));
    return EasyMock.expectLastCall();
}

public static AndroidLibrary createMockLibrary(String allResources, String publicResources,
        List<AndroidLibrary> dependencies)
        throws IOException {
    final File tempDir = TestUtils.createTempDirDeletedOnExit();

    Files.write(allResources, new File(tempDir, FN_RESOURCE_TEXT), Charsets.UTF_8);
    File publicTxtFile = new File(tempDir, FN_PUBLIC_TXT);
    if (publicResources != null) {
        Files.write(publicResources, publicTxtFile, Charsets.UTF_8);
    }
    AndroidLibrary library = createNiceMock(AndroidLibrary.class);
    expect(library.getPublicResources()).andReturn(publicTxtFile).anyTimes();

    // Work around wildcard capture
    //expect(mock.getLibraryDependencies()).andReturn(dependencies).anyTimes();
    IExpectationSetters setter = expect(library.getLibraryDependencies());
    //noinspection unchecked
    setter.andReturn(dependencies);
    setter.anyTimes();

    replay(library);
    return library;
}

/**
 * Used to specify expectations on private methods. Use this method to
 * handle overloaded methods.
 */
@SuppressWarnings("all")
public static synchronized <T> IExpectationSetters<T> expectPrivate(Object instance, String methodName,
                                                                    Class<?>[] parameterTypes, Object... arguments) throws Exception {

    if (arguments == null) {
        arguments = new Object[0];
    }

    if (instance == null) {
        throw new IllegalArgumentException("instance cannot be null.");
    } else if (arguments.length != parameterTypes.length) {
        throw new IllegalArgumentException(
                "The length of the arguments must be equal to the number of parameter types.");
    }

    Method foundMethod = Whitebox.getMethod(instance.getClass(), methodName, parameterTypes);

    WhiteboxImpl.throwExceptionIfMethodWasNotFound(instance.getClass(), methodName, foundMethod, parameterTypes);

    return doExpectPrivate(instance, foundMethod, arguments);
}

/**
 * Used to specify expectations on methods using the method name at a
 * specific place in the class hierarchy (specified by the
 * <code>where</code> parameter). Works on for example private or package
 * private methods.
 * <p>
 * Use this for overloaded methods.
 */
public static synchronized <T> IExpectationSetters<T> expectPrivate(Object instance, String methodName,
                                                                    Class<?> where, Class<?>[] parameterTypes, Object... arguments) throws Exception {
    if (instance == null) {
        throw new IllegalArgumentException("Instance or class to expect cannot be null.");
    }
    Method[] methods = null;
    if (methodName != null) {
        if (parameterTypes == null) {
            methods = Whitebox.getMethods(where, methodName);
        } else {
            methods = new Method[] { Whitebox.getMethod(where, methodName, parameterTypes) };
        }
    }
    Method methodToExpect;
    if (methods != null && methods.length == 1) {
        methodToExpect = methods[0];
    } else {
        methodToExpect = WhiteboxImpl.findMethodOrThrowException(instance, null, methodName, arguments);
    }

    return doExpectPrivate(instance, methodToExpect, arguments);
}

private void expectConvertKeyValue(boolean anyTimes) {
    IExpectationSetters<byte[]> convertKeyExpect = EasyMock.expect(keyConverter.fromConnectData(TOPIC, KEY_SCHEMA, KEY));
    if (anyTimes)
        convertKeyExpect.andStubReturn(SERIALIZED_KEY);
    else
        convertKeyExpect.andReturn(SERIALIZED_KEY);
    IExpectationSetters<byte[]> convertValueExpect = EasyMock.expect(valueConverter.fromConnectData(TOPIC, RECORD_SCHEMA, RECORD));
    if (anyTimes)
        convertValueExpect.andStubReturn(SERIALIZED_RECORD);
    else
        convertValueExpect.andReturn(SERIALIZED_RECORD);
}

private void expectTaskCommitRecord(boolean anyTimes, boolean succeed) throws InterruptedException {
    sourceTask.commitRecord(EasyMock.anyObject(SourceRecord.class));
    IExpectationSetters<Void> expect = EasyMock.expectLastCall();
    if (!succeed) {
        expect = expect.andThrow(new RuntimeException("Error committing record in source task"));
    }
    if (anyTimes) {
        expect.anyTimes();
    }
}

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;
}

private IExpectationSetters<Boolean> expectAssigned(
    IScheduledTask task,
    Map<String, TaskGroupKey> reservationMap) {

  return expect(assigner.maybeAssign(
      storageUtil.mutableStoreProvider,
      new ResourceRequest(task.getAssignedTask().getTask(), EMPTY),
      TaskGroupKey.from(task.getAssignedTask().getTask()),
      Tasks.id(task),
      reservationMap));
}

private IExpectationSetters<?> expectGetHostRack(String host, String rackToReturn) {
  IHostAttributes attributes = IHostAttributes.build(new HostAttributes()
      .setHost(host)
      .setAttributes(ImmutableSet.of(
          new Attribute().setName("rack").setValues(ImmutableSet.of(rackToReturn)))));
  return expect(storageUtil.attributeStore.getHostAttributes(host))
      .andReturn(Optional.of(attributes));
}

private IExpectationSetters<?> expectFetch(IAssignedTask... results) {
  ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder();
  for (IAssignedTask result : results) {
    tasks.add(IScheduledTask.build(new ScheduledTask().setAssignedTask(result.newBuilder())));
  }

  return expect(store.fetchTasks(Query.jobScoped(JOB).active()))
      .andReturn(tasks.build());
}

private IExpectationSetters<Set<SchedulingFilter.Veto>> expectFiltering(
    final Optional<Veto> veto) {

  return expect(schedulingFilter.filter(
      EasyMock.anyObject(),
      EasyMock.anyObject()))
      .andAnswer(
          veto::asSet);
}

private IExpectationSetters<Position> expectWrite(String content) throws Exception {
  return expect(
      logWriter.append(EasyMock.aryEq(content.getBytes(StandardCharsets.UTF_8)),
          // Cast is needed to prevent NullPointerException on unboxing.
          EasyMock.eq((long) WRITE_TIMEOUT.getValue()),
          EasyMock.eq(WRITE_TIMEOUT.getUnit().getTimeUnit())));
}

private IExpectationSetters<List<Log.Entry>> expectRead(Position position) throws Exception {
  expectSetPosition(position);
  return expect(logReader.read(
      position,
      position,
      READ_TIMEOUT.getValue(),
      READ_TIMEOUT.getUnit().getTimeUnit()));
}

private IExpectationSetters<Object> expectShiroAfterAuthFilter()
    throws ServletException, IOException {

  shiroAfterAuthFilter.doFilter(
      isA(HttpServletRequest.class),
      isA(HttpServletResponse.class),
      isA(FilterChain.class));

  return expectLastCall().andAnswer(() -> {
    Object[] args = getCurrentArguments();
    ((FilterChain) args[2]).doFilter((HttpServletRequest) args[0], (HttpServletResponse) args[1]);
    return null;
  });
}

private IExpectationSetters<?> expectInstanceQuotaCheck(
    ITaskConfig config,
    QuotaCheckResult result) {

  return expect(quotaManager.checkInstanceAddition(
      config,
      1,
      storageUtil.mutableStoreProvider)).andReturn(result);
}

private IExpectationSetters<?> expectCronQuotaCheck(
    IJobConfiguration config,
    QuotaCheckResult result) {

  return expect(quotaManager.checkCronUpdate(config, storageUtil.mutableStoreProvider))
      .andReturn(result);
}

private IExpectationSetters<?> expectCronJobs(IJobConfiguration... jobs) {
  ImmutableSet.Builder<IJobConfiguration> builder = ImmutableSet.builder();
  for (IJobConfiguration job : jobs) {
    builder.add(job);
  }

  return expect(storageUtil.jobStore.fetchJobs()).andReturn(builder.build());
}

/**
 * Sets an expectation for a snapshot.
 *
 * @param snapshot Expected snapshot.
 * @return An expectation setter.
 */
public IExpectationSetters<Position> expectSnapshot(DeduplicatedSnapshot snapshot) {
  try {
    LogEntry entry = Entries.deflate(LogEntry.deduplicatedSnapshot(snapshot));
    return expect(stream.append(sameEntry(entry)));
  } catch (CodingException e) {
    throw Throwables.propagate(e);
  }
}

private IExpectationSetters<?> expectGateClosed() throws Exception {
  return expect(gatedWorkQueue.closeDuring(EasyMock.anyObject()))
      .andAnswer(() -> {
        GatedOperation<?, ?> op = (GatedOperation<?, ?>) EasyMock.getCurrentArguments()[0];
        return op.doWithGateClosed();
      });
}

private IExpectationSetters<CloseableHttpResponse> backendExpectsAnyRequestAndReturn(
    final HttpResponse response) throws Exception {
    final CloseableHttpResponse resp = mockBackend.execute(EasyMock.isA(HttpRoute.class),
        EasyMock.isA(HttpRequestWrapper.class), EasyMock.isA(HttpClientContext.class),
        EasyMock.<HttpExecutionAware> isNull());
    return EasyMock.expect(resp).andReturn(Proxies.enhanceResponse(response));
}

protected IExpectationSetters<CloseableHttpResponse> backendExpectsRequestAndReturn(
    final HttpRequestWrapper request, final HttpResponse response) throws Exception {
    final CloseableHttpResponse resp = mockBackend.execute(EasyMock.isA(HttpRoute.class),
        EasyMock.eq(request), EasyMock.isA(HttpClientContext.class),
        EasyMock.<HttpExecutionAware> isNull());
    return EasyMock.expect(resp).andReturn(Proxies.enhanceResponse(response));
}

protected IExpectationSetters<CloseableHttpResponse> backendExpectsRequestAndReturn(
    final HttpRequestWrapper request, final CloseableHttpResponse response) throws Exception {
    final CloseableHttpResponse resp = mockBackend.execute(EasyMock.isA(HttpRoute.class),
        EasyMock.eq(request), EasyMock.isA(HttpClientContext.class),
        EasyMock.<HttpExecutionAware> isNull());
    return EasyMock.expect(resp).andReturn(response);
}

protected IExpectationSetters<CloseableHttpResponse> backendExpectsAnyRequestAndThrows(
    final Throwable throwable) throws Exception {
    final CloseableHttpResponse resp = mockBackend.execute(EasyMock.isA(HttpRoute.class),
        EasyMock.isA(HttpRequestWrapper.class), EasyMock.isA(HttpClientContext.class),
        EasyMock.<HttpExecutionAware> isNull());
    return EasyMock.expect(resp).andThrow(throwable);
}

protected IExpectationSetters<CloseableHttpResponse> backendCaptureRequestAndReturn(
    final Capture<HttpRequestWrapper> cap, final HttpResponse response) throws Exception {
    final CloseableHttpResponse resp = mockBackend.execute(EasyMock.isA(HttpRoute.class),
        EasyMock.capture(cap), EasyMock.isA(HttpClientContext.class),
        EasyMock.<HttpExecutionAware> isNull());
    return EasyMock.expect(resp).andReturn(Proxies.enhanceResponse(response));
}

private IExpectationSetters<CloseableHttpResponse> implExpectsAnyRequestAndReturn(
        final CloseableHttpResponse response) throws Exception {
    final CloseableHttpResponse resp = impl.callBackend(
            EasyMock.isA(HttpRoute.class),
            EasyMock.isA(HttpRequestWrapper.class),
            EasyMock.isA(HttpClientContext.class),
            EasyMock.<HttpExecutionAware>isNull());
    return EasyMock.expect(resp).andReturn(response);
}

protected IExpectationSetters<CloseableHttpResponse> backendExpectsAnyRequest() throws Exception {
    final CloseableHttpResponse resp = mockBackend.execute(
            EasyMock.isA(HttpRoute.class),
            EasyMock.isA(HttpRequestWrapper.class),
            EasyMock.isA(HttpClientContext.class),
            EasyMock.<HttpExecutionAware>isNull());
    return EasyMock.expect(resp);
}

protected IExpectationSetters<CloseableHttpResponse> backendExpectsAnyRequestAndReturn(
        final HttpResponse reponse) throws Exception {
    final CloseableHttpResponse resp = mockBackend.execute(
            EasyMock.isA(HttpRoute.class),
            EasyMock.isA(HttpRequestWrapper.class),
            EasyMock.isA(HttpClientContext.class),
            EasyMock.<HttpExecutionAware>isNull());
    return EasyMock.expect(resp).andReturn(Proxies.enhanceResponse(reponse));
}

/**
 * Used to specify expectations on methods using the method name. Works on
 * for example private or package private methods.
 */
public static synchronized <T> IExpectationSetters<T> expectPrivate(Object instance, String methodName,
                                                                    Object... arguments) throws Exception {
    if (instance == null) {
        throw new IllegalArgumentException("Instance or class cannot be null.");
    }

    return expectPrivate(instance, methodName, Whitebox.getType(instance), arguments);
}

@SuppressWarnings("unchecked")
private static <T> IExpectationSetters<T> doExpectNew(Class<T> type, MockStrategy mockStrategy,
                                                      Class<?>[] parameterTypes, Object... arguments) throws Exception {
    if (type == null) {
        throw new IllegalArgumentException("type cannot be null");
    } else if (mockStrategy == null) {
        throw new IllegalArgumentException("Internal error: Mock strategy cannot be null");
    }

    final boolean isNiceMock = mockStrategy instanceof NiceMockStrategy;

    final Class<T> unmockedType = (Class<T>) WhiteboxImpl.getUnmockedType(type);
    if (!isNiceMock) {
        if (parameterTypes == null) {
            WhiteboxImpl.findUniqueConstructorOrThrowException(type, arguments);
        } else {
            WhiteboxImpl.getConstructor(unmockedType, parameterTypes);
        }
    }

    /*
       * Check if this type has been mocked before
       */
    NewInvocationControl<IExpectationSetters<T>> newInvocationControl = (NewInvocationControl<IExpectationSetters<T>>) MockRepository
            .getNewInstanceControl(unmockedType);
    if (newInvocationControl == null) {
        InvocationSubstitute<T> mock = doMock(InvocationSubstitute.class, false, mockStrategy, null,
                (Method[]) null);
        newInvocationControl = new NewInvocationControlImpl<T>(mock, type);
        MockRepository.putNewInstanceControl(type, newInvocationControl);
        MockRepository.addObjectsToAutomaticallyReplayAndVerify(WhiteboxImpl.getUnmockedType(type));
    }

    if (isNiceMock && (arguments == null || arguments.length == 0)) {
        return null;
    }
    return newInvocationControl.expectSubstitutionLogic(arguments);
}

private MockUnit.Block executeCallback(String defaultUrl, String logoutUrlPattern,
    String matchers, boolean multiProfile, Throwable x) {
  return unit -> {
    SecurityLogic action = unit.get(SecurityLogic.class);

    IExpectationSetters<Object> expect = expect(
        action.perform(eq(unit.get(WebContext.class)), eq(unit.get(Config.class)),
            isA(Pac4jGrantAccessAdapter.class), eq(unit.get(HttpActionAdapter.class)),
            eq(defaultUrl), eq(logoutUrlPattern), eq(matchers), eq(multiProfile)));
    if (x == null)
      expect.andReturn(null);
    else
      expect.andThrow(x);
  };
}

protected SolrQueryRequest prepareCollAndUser(SolrCore core, SolrQueryRequest request,
    String collection, String user, boolean onlyOnce) throws Exception {
  CloudDescriptor mCloudDescriptor = EasyMock.createMock(CloudDescriptor.class);
  IExpectationSetters getCollNameExpect = EasyMock.expect(mCloudDescriptor.getCollectionName()).andReturn(collection);
  getCollNameExpect.anyTimes();
  IExpectationSetters getShardIdExpect = EasyMock.expect(mCloudDescriptor.getShardId()).andReturn("shard1");
  getShardIdExpect.anyTimes();
  EasyMock.replay(mCloudDescriptor);
  CoreDescriptor coreDescriptor = core.getCoreDescriptor();
  Field cloudDescField = CoreDescriptor.class.getDeclaredField("cloudDesc");
  cloudDescField.setAccessible(true);
  cloudDescField.set(coreDescriptor, mCloudDescriptor);

  HttpServletRequest httpServletRequest = EasyMock.createMock(HttpServletRequest.class);
  IExpectationSetters getAttributeUserExpect =
      EasyMock.expect(httpServletRequest.getAttribute(USER_NAME)).andReturn(user);
  if (!onlyOnce) {
    getAttributeUserExpect.anyTimes();
  }
  IExpectationSetters getAttributeDoAsUserExpect =
      EasyMock.expect(httpServletRequest.getAttribute(DO_AS_USER_NAME)).andReturn(null);
  if (!onlyOnce) {
    getAttributeDoAsUserExpect.anyTimes();
  }
  EasyMock.replay(httpServletRequest);
  request.getContext().put("httpRequest", httpServletRequest);
  return request;
}

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;
}

@SuppressWarnings("unchecked")
private IExpectationSetters<Object> expectRebalanceDuringPoll() throws Exception {
    final List<TopicPartition> partitions = Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3);

    final long startOffset = 40L;
    final Map<TopicPartition, Long> offsets = new HashMap<>();
    offsets.put(TOPIC_PARTITION, startOffset);

    EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
            new IAnswer<ConsumerRecords<byte[], byte[]>>() {
                @Override
                public ConsumerRecords<byte[], byte[]> answer() throws Throwable {
                    // "Sleep" so time will progress
                    time.sleep(1L);

                    sinkTaskContext.getValue().offset(offsets);
                    rebalanceListener.getValue().onPartitionsAssigned(partitions);

                    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(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET);
    EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(FIRST_OFFSET);
    EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(FIRST_OFFSET);

    sinkTask.open(partitions);
    EasyMock.expectLastCall();

    consumer.seek(TOPIC_PARTITION, startOffset);
    EasyMock.expectLastCall();

    EasyMock.expect(keyConverter.toConnectData(TOPIC, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY));
    EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE));
    sinkTask.put(EasyMock.anyObject(Collection.class));
    return EasyMock.expectLastCall();
}

private IExpectationSetters<ServerSet.EndpointStatus> expectJoin() throws Exception {
  return expect(serverSet.join(PRIMARY_ENDPOINT, AUX_ENDPOINTS));
}

private IExpectationSetters<?> expectGetReservation(IScheduledTask task, String slaveId) {
  return expect(reservations.getByValue(TaskGroupKey.from(task.getAssignedTask().getTask())))
      .andReturn(ImmutableSet.of(slaveId));
}

private IExpectationSetters<?> expectNoReservation(IScheduledTask task) {
  return expect(reservations.getByValue(TaskGroupKey.from(task.getAssignedTask().getTask())))
      .andReturn(ImmutableSet.of());
}

private IExpectationSetters<?> expectAsMap(Map<String, TaskGroupKey> map) {
  return expect(reservations.asMap()).andReturn(map);
}

private IExpectationSetters<String> expectTaskKilled() {
  driver.killTask(EasyMock.anyObject());
  return expectLastCall();
}