Maintenance { @GET @Produces(MediaType.APPLICATION_JSON) public Response getHosts() { return storage.read(storeProvider -> { Multimap<MaintenanceMode, String> hostsByMode = Multimaps.transformValues( Multimaps.index( storeProvider.getAttributeStore().getHostAttributes(), IHostAttributes::getMode), IHostAttributes::getHost); Map<MaintenanceMode, Object> hosts = ImmutableMap.of( DRAINED, ImmutableSet.copyOf(hostsByMode.get(DRAINED)), SCHEDULED, ImmutableSet.copyOf(hostsByMode.get(SCHEDULED)), DRAINING, getTasksByHosts(storeProvider, hostsByMode.get(DRAINING)).asMap()); return Response.ok(hosts).build(); }); } @Inject  Maintenance(Storage storage); @GET @Produces(MediaType.APPLICATION_JSON) Response getHosts();  }
@Test public void testNoDrainingHosts() { storage.write((Quiet) storeProvider -> { ScheduledTask pending = TaskTestUtil.makeTask("a", TaskTestUtil.JOB) .newBuilder() .setStatus(ScheduleStatus.PENDING); storeProvider.getAttributeStore().saveHostAttributes(IHostAttributes.build( new HostAttributes() .setHost("b") .setSlaveId("b") .setMode(MaintenanceMode.NONE) )); ScheduledTask assigned = TaskTestUtil.makeTask("b", TaskTestUtil.JOB) .newBuilder() .setStatus(ScheduleStatus.ASSIGNED); assigned.getAssignedTask() .setInstanceId(0) .setSlaveHost("b") .setSlaveId("b"); storeProvider.getUnsafeTaskStore().saveTasks( IScheduledTask.setFromBuilders(ImmutableList.of(pending, assigned))); }); @SuppressWarnings("unchecked") Map<MaintenanceMode, Object> result = (Map<MaintenanceMode, Object>) maintenance.getHosts().getEntity(); assertEquals(ImmutableSet.of(), result.get(SCHEDULED)); assertEquals(ImmutableSet.of(), result.get(DRAINED)); assertEquals(ImmutableMap.of(), result.get(DRAINING)); }