ElasticsearchScheduler implements Scheduler { @Override public void resourceOffers(SchedulerDriver driver, List<Protos.Offer> offers) { if (!frameworkState.isRegistered()) { LOGGER.debug("Not registered, can't accept resource offers."); return; } for (Protos.Offer offer : offers) { final OfferStrategy.OfferResult result = offerStrategy.evaluate(offer); if (!result.acceptable) { LOGGER.debug("Declined offer: " + flattenProtobufString(offer.toString()) + "Reason: " + result.reason.orElse("Unknown")); driver.declineOffer(offer.getId()); } else { Protos.TaskInfo taskInfo = taskInfoFactory.createTask(configuration, frameworkState, offer, new Clock()); LOGGER.debug(taskInfo.toString()); driver.launchTasks(Collections.singleton(offer.getId()), Collections.singleton(taskInfo)); ESTaskStatus esTask = new ESTaskStatus(zookeeperStateDriver, frameworkState.getFrameworkID(), taskInfo, new StatePath(zookeeperStateDriver)); clusterState.addTask(esTask); frameworkState.announceNewTask(esTask); } } } ElasticsearchScheduler(Configuration configuration, FrameworkState frameworkState, ClusterState clusterState, TaskInfoFactory taskInfoFactory, OfferStrategy offerStrategy, SerializableState zookeeperStateDriver); Map<String, Task> getTasks(); void run(SchedulerDriver schedulerDriver); void reapTasks(); @Override void registered(SchedulerDriver driver, Protos.FrameworkID frameworkId, Protos.MasterInfo masterInfo); @Override void reregistered(SchedulerDriver driver, Protos.MasterInfo masterInfo); @Override void resourceOffers(SchedulerDriver driver, List<Protos.Offer> offers); @Override void offerRescinded(SchedulerDriver driver, Protos.OfferID offerId); @Override void statusUpdate(SchedulerDriver driver, Protos.TaskStatus status); @Override void frameworkMessage(SchedulerDriver driver, Protos.ExecutorID executorId, Protos.SlaveID slaveId, byte[] data); @Override void disconnected(SchedulerDriver driver); @Override void slaveLost(SchedulerDriver driver, Protos.SlaveID slaveId); @Override void executorLost(SchedulerDriver driver, Protos.ExecutorID executorId, Protos.SlaveID slaveId, int status); @Override void error(SchedulerDriver driver, String message); void shutdown(SchedulerDriver driver);  }
@Test public void willDeclineOfferIfStrategyDeclinesOffer() { Protos.Offer offer = newOffer("host1").build(); when(offerStrategy.evaluate(offer)).thenReturn(OfferStrategy.OfferResult.decline("Test")); when(frameworkState.isRegistered()).thenReturn(true); scheduler.resourceOffers(driver, singletonList(offer)); verify(driver).declineOffer(offer.getId()); }
@Test public void testResourceOffers_launchTasks() { final Protos.Offer offer = newOffer("host3").build(); when(offerStrategy.evaluate(offer)).thenReturn(OfferStrategy.OfferResult.accept()); when(frameworkState.isRegistered()).thenReturn(true); Protos.TaskInfo taskInfo = ProtoTestUtil.getDefaultTaskInfo(); when(taskInfoFactory.createTask(any(), any(), any(), any())).thenReturn(taskInfo); scheduler.resourceOffers(driver, singletonList(offer)); verify(driver).launchTasks(singleton(offer.getId()), singleton(taskInfo)); }