public static void setPermanentlyFailed(StateStore stateStore, Collection<Protos.TaskInfo> failedTasks) { stateStore.storeTasks(failedTasks.stream() .map(taskInfo -> taskInfo.toBuilder() .setLabels(new TaskLabelWriter(taskInfo).setPermanentlyFailed().toProto()) .build()) .collect(Collectors.toList())); } }
@Test public void testPermanentlyFailedTaskNeedsRecovery() throws Exception { ConfigStore<ServiceSpec> configStore = newConfigStore(persister); Protos.TaskInfo taskInfo = newTaskInfo("name-0-node", configStore); // Set status as RUNNING Protos.TaskStatus taskStatus = StateStoreUtilsTest.newTaskStatus(taskInfo, Protos.TaskState.TASK_RUNNING); // Mark task as permanently failed taskInfo = taskInfo.toBuilder() .setLabels(new TaskLabelWriter(taskInfo).setPermanentlyFailed().toProto()) .build(); // Even though the TaskStatus is RUNNING, it can now be recovered since it has been marked as // permanently failed. assertThat(TaskUtils.getTasksNeedingRecovery(configStore, Collections.singleton(taskInfo), Collections.singleton(taskStatus)), is(ImmutableList.of(taskInfo))); }
@Test public void getRegionTaskKeys() { MaxPerRegionRule rule = new MaxPerRegionRule(2); Protos.TaskInfo taskInfo = TestConstants.TASK_INFO.toBuilder() .setLabels( new TaskLabelWriter(TestConstants.TASK_INFO) .setRegion(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getRegion()) .toProto()) .build(); Collection<String> keys = rule.getKeys(taskInfo); Assert.assertEquals(1, keys.size()); Assert.assertEquals( TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getRegion().getName(), keys.stream().findFirst().get()); } }
@Test public void testReadinessCheckTagging() throws TaskException { Protos.HealthCheck inReadinessCheck = Protos.HealthCheck.newBuilder() .setDelaySeconds(1.0) .build(); Protos.TaskInfo.Builder builder = getTestTaskInfo().toBuilder(); builder.setLabels(new TaskLabelWriter(builder) .setReadinessCheck(inReadinessCheck) .toProto()); Protos.HealthCheck outReadinessCheck = new TaskLabelWriter(builder.build()) { @Override public Optional<HealthCheck> getReadinessCheck() throws TaskException { return super.getReadinessCheck(); } }.getReadinessCheck().get(); Assert.assertEquals(inReadinessCheck.getDelaySeconds(), outReadinessCheck.getDelaySeconds(), 0.0); }
private static Protos.TaskInfo newTaskInfo( final String taskName, final ConfigStore<ServiceSpec> configStore) throws ConfigStoreException { Protos.TaskInfo.Builder taskInfoBuilder = newTaskInfo(taskName).toBuilder(); // POD type final UUID targetConfig = configStore.getTargetConfig(); final int podIndex = 0; final String podType = configStore.fetch(targetConfig).getPods().get(podIndex).getType(); // create default labels: taskInfoBuilder.setLabels(new TaskLabelWriter(taskInfoBuilder) .setTargetConfiguration(targetConfig) .setType(podType) .setIndex(podIndex) .toProto()); return taskInfoBuilder.build(); }
@Test public void getRegionTaskKey() { RoundRobinByRegionRule rule = new RoundRobinByRegionRule(2); Protos.TaskInfo taskInfo = TestConstants.TASK_INFO.toBuilder() .setLabels( new TaskLabelWriter(TestConstants.TASK_INFO) .setRegion(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getRegion()) .toProto()) .build(); String key = rule.getKey(taskInfo); Assert.assertNotNull(key); Assert.assertEquals(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getRegion().getName(), key); } }
@Test public void getZoneTaskKey() { RoundRobinByZoneRule rule = new RoundRobinByZoneRule(2); Protos.TaskInfo taskInfo = TestConstants.TASK_INFO.toBuilder() .setLabels( new TaskLabelWriter(TestConstants.TASK_INFO) .setZone(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getZone()) .toProto()) .build(); String key = rule.getKey(taskInfo); Assert.assertNotNull(key); Assert.assertEquals(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getZone().getName(), key); } }
@Test public void testUpdateResourcesNoHarmForAcceptableTaskInfo() throws TaskException { Protos.TaskInfo withLabel = baseTaskInfo.toBuilder() .setLabels(new TaskLabelWriter(baseTaskInfo).setType(TestConstants.TASK_TYPE).toProto()) .build(); Assert.assertFalse(persistentLaunchRecorder.getPodInstance(withLabel).isPresent()); }
private static TaskInfo getTaskInfo(String name, String host) { TaskInfo.Builder infoBuilder = TaskTestUtils.getTaskInfo(Collections.emptyList()).toBuilder() .setName(name) .setTaskId(CommonIdUtils.toTaskId(TestConstants.SERVICE_NAME, name)); infoBuilder.setLabels(new TaskLabelWriter(infoBuilder).setHostname(offerWithHost(host)).toProto()); return infoBuilder.build(); }
private static TaskInfo getTask(String id, Offer offer) { TaskInfo.Builder taskBuilder = TaskTestUtils.getTaskInfo(Collections.emptyList()).toBuilder(); taskBuilder.getTaskIdBuilder().setValue(id); try { taskBuilder.setName(CommonIdUtils.toTaskName(taskBuilder.getTaskId())); } catch (Exception e) { throw new IllegalStateException(e); } taskBuilder.setLabels(new TaskLabelWriter(taskBuilder).setOfferAttributes(offer).toProto()); return taskBuilder.build(); }
private static Protos.TaskInfo getTask(String id, Protos.Offer offer) { Protos.TaskInfo.Builder taskBuilder = TaskTestUtils.getTaskInfo(Collections.emptyList()).toBuilder(); taskBuilder.getTaskIdBuilder().setValue(id); try { taskBuilder.setName(CommonIdUtils.toTaskName(taskBuilder.getTaskId())); } catch (Exception e) { throw new IllegalStateException(e); } taskBuilder.setLabels(new TaskLabelWriter(taskBuilder).setHostname(offer).toProto()); return taskBuilder.build(); }
public static Protos.TaskInfo withFailedFlag(Protos.TaskInfo task) { return task.toBuilder().setLabels(new TaskLabelWriter(task).setPermanentlyFailed().toProto()).build(); } }
@Test public void testAdditionalLabels() throws TaskException { Map<String, String> labels = new HashMap<String, String>(); labels.put("label1", "label1-value"); labels.put("label2", "label2-value"); Protos.TaskInfo.Builder builder = getTestTaskInfo().toBuilder(); builder.setLabels(new TaskLabelWriter(builder).setAdditionalLabels(labels).toProto()); Assert.assertEquals("label1-value", new LabelReader("", builder.getLabels()).getOrThrow("label1")); Assert.assertEquals("label2-value", new LabelReader("", builder.getLabels()).getOrThrow("label2")); }
@Test public void testReadWriteZone() { Assert.assertFalse(new TaskLabelReader(getTestTaskInfo()).getZone().isPresent()); Protos.TaskInfo.Builder tb = getTestTaskInfo().toBuilder(); tb.setLabels(new TaskLabelWriter(tb).setZone(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getZone()).toProto()); Assert.assertTrue(new TaskLabelReader(tb.build()).getZone().isPresent()); Assert.assertEquals(TestConstants.ZONE, new TaskLabelReader(tb.build()).getZone().get()); }
@Before public void beforeEach() { MockitoAnnotations.initMocks(this); podInstance = PodTestUtils.getPodInstance(0); taskInfo = TestConstants.TASK_INFO.toBuilder().setLabels( new TaskLabelWriter(TestConstants.TASK_INFO) .setType("different-type") .setIndex(100) .toProto()) .build(); offer = OfferTestUtils.getEmptyOfferBuilder().build(); }
private static TaskInfo getTaskInfo(String taskName, String attrName, String attrVal) { TaskInfo.Builder infoBuilder = TaskTestUtils.getTaskInfo(Collections.emptyList()).toBuilder() .setName(taskName) .setTaskId(CommonIdUtils.toTaskId(TestConstants.SERVICE_NAME, taskName)); infoBuilder.setLabels(new TaskLabelWriter(infoBuilder) .setOfferAttributes(offerWithAttribute(attrName, attrVal)) .toProto()); return infoBuilder.build(); }
@Test public void testReadWriteRegion() { Assert.assertFalse(new TaskLabelReader(getTestTaskInfo()).getRegion().isPresent()); Protos.TaskInfo.Builder tb = getTestTaskInfo().toBuilder(); tb.setLabels(new TaskLabelWriter(tb).setRegion(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getRegion()).toProto()); Assert.assertTrue(new TaskLabelReader(tb.build()).getRegion().isPresent()); Assert.assertEquals(TestConstants.LOCAL_REGION, new TaskLabelReader(tb.build()).getRegion().get()); }
@Test public void testSetTargetConfiguration() throws Exception { Protos.TaskInfo.Builder taskBuilder = getTestTaskInfo().toBuilder(); taskBuilder.setLabels(new TaskLabelWriter(taskBuilder) .setTargetConfiguration(testTargetConfigurationId) .toProto()); Assert.assertEquals(testTargetConfigurationId, new TaskLabelReader(taskBuilder).getTargetConfiguration()); }
private static TaskInfo getTask(String type, String id, String agent) { TaskInfo.Builder taskBuilder = TaskTestUtils.getTaskInfo(Collections.emptyList()).toBuilder(); taskBuilder.getTaskIdBuilder().setValue(id); try { taskBuilder.setName(CommonIdUtils.toTaskName(taskBuilder.getTaskId())); } catch (Exception e) { throw new IllegalStateException(e); } taskBuilder.getSlaveIdBuilder().setValue(agent); taskBuilder.setLabels(new TaskLabelWriter(taskBuilder).setType(type).toProto()); return taskBuilder.build(); }
@Test public void testSetGetTaskType() throws TaskException { Protos.TaskInfo.Builder builder = getTestTaskInfo().toBuilder(); builder.setLabels(new TaskLabelWriter(builder).setType("foo").toProto()); Assert.assertEquals("foo", new TaskLabelReader(builder).getType()); builder = getTestTaskInfo().toBuilder(); builder.setLabels(new TaskLabelWriter(builder).setType("").toProto()); Assert.assertEquals("", new TaskLabelReader(builder).getType()); }