Expect.java 文件源码

/**
 * Verifies that the resources for the provided task names have been unreserved.
 */
public static Expect unreservedTasks(Collection<String> taskNames) {
    return new Expect() {
        // Use this form instead of using ArgumentCaptor.forClass() to avoid problems with typecasting generics:
        @Captor private ArgumentCaptor<Collection<Protos.OfferID>> offerIdsCaptor;
        @Captor private ArgumentCaptor<Collection<Protos.Offer.Operation>> operationsCaptor;

        @Override
        public void expect(ClusterState state, SchedulerDriver mockDriver) {
            MockitoAnnotations.initMocks(this);
            verify(mockDriver, atLeastOnce())
                    .acceptOffers(offerIdsCaptor.capture(), operationsCaptor.capture(), any());
            Assert.assertEquals(state.getLastOffer().getId(), offerIdsCaptor.getValue().iterator().next());
            Collection<String> expectedResourceIds = taskNames.stream()
                    .map(taskName ->
                            ResourceUtils.getResourceIds(state.getLastLaunchedTask(taskName).getResourcesList()))
                    .flatMap(List::stream)
                    .collect(Collectors.toList());
            Assert.assertFalse(String.format("Expected some resource ids for tasks: %s, got none", taskNames),
                    expectedResourceIds.isEmpty());
            Collection<String> unreservedResourceIds = new ArrayList<>();
            for (Protos.Offer.Operation operation : operationsCaptor.getValue()) {
                if (operation.getType().equals(Protos.Offer.Operation.Type.DESTROY)) {
                    // Destroy volume(s)
                    unreservedResourceIds.addAll(
                            ResourceUtils.getResourceIds(operation.getDestroy().getVolumesList()));
                } else if (operation.getType().equals(Protos.Offer.Operation.Type.UNRESERVE)) {
                    // Unreserve resource(s)
                    unreservedResourceIds.addAll(
                            ResourceUtils.getResourceIds(operation.getUnreserve().getResourcesList()));
                }
            }
            Assert.assertTrue(
                    String.format("Expected unreserved resource ids: %s, got ids: %s",
                            expectedResourceIds, unreservedResourceIds),
                    unreservedResourceIds.containsAll(expectedResourceIds)
                    && expectedResourceIds.containsAll(unreservedResourceIds));
        }

        @Override
        public String getDescription() {
            return String.format("Resources for tasks have been unreserved: %s", taskNames);
        }
    };
}