FrameworkState { public void markRegistered(Protos.FrameworkID frameworkId, SchedulerDriver driver) { if (!registered.compareAndSet(false, true)) { throw new IllegalStateException("Framework can not be marked as registered twice"); } try { statePath.mkdir(FRAMEWORKID_KEY); zookeeperStateDriver.set(FRAMEWORKID_KEY, frameworkId); LOGGER.debug("FrameworkID stored in zookeeper: " + FRAMEWORKID_KEY + " = " + frameworkId); } catch (IOException e) { LOGGER.error("Unable to store framework ID in zookeeper", e); } this.driver = driver; final ClusterState clusterState = new ClusterState(zookeeperStateDriver, this); registeredListeners.forEach(listener -> listener.accept(clusterState)); } FrameworkState(SerializableState zookeeperStateDriver); Protos.FrameworkID getFrameworkID(); void markRegistered(Protos.FrameworkID frameworkId, SchedulerDriver driver); void destroy(); SchedulerDriver getDriver(); void onRegistered(Consumer<ClusterState> listener); void onNewTask(Consumer<ESTaskStatus> listener); boolean isRegistered(); void announceNewTask(ESTaskStatus esTask); void announceStatusUpdate(Protos.TaskStatus taskStatus); void onStatusUpdate(Consumer<Protos.TaskStatus> listener); static final Protos.FrameworkID EMPTY_ID; }
@Test public void testSetFrameworkID() throws IOException { frameworkState.markRegistered(FRAMEWORK_ID, driver); verify(state, times(1)).set(anyString(), eq(FRAMEWORK_ID)); }
@Test public void testHandleSetException() throws IOException { doThrow(IOException.class).when(state).set(anyString(), any()); frameworkState.markRegistered(FRAMEWORK_ID, driver); }