use of com.netflix.titus.master.integration.v3.scenario.TaskScenarioBuilder in project titus-control-plane by Netflix.
the class JobIpAllocationsTest method testWaitingTaskContext.
/**
* Tests a job waiting for an in use IP allocation has updated task context fields.
*/
@Test(timeout = 30_000)
// TODO Read static IP allocation status from pod message field, and add it to the task context.
@Ignore
public void testWaitingTaskContext() throws Exception {
JobDescriptor<ServiceJobExt> firstIpJobDescriptor = ONE_TASK_SERVICE_JOB;
JobDescriptor<ServiceJobExt> secondIpJobDescriptor = firstIpJobDescriptor.but(j -> j.getJobGroupInfo().toBuilder().withSequence("v001"));
// Schedule the first task and ensure it's in the correct zone with the correct task context
jobsScenarioBuilder.schedule(firstIpJobDescriptor, jobScenarioBuilder -> jobScenarioBuilder.template(ScenarioTemplates.startTasksInNewJob()).allTasks(taskScenarioBuilder -> taskScenarioBuilder.expectTaskContext(TaskAttributes.TASK_ATTRIBUTES_IP_ALLOCATION_ID, getIpAllocationIdFromJob(0, firstIpJobDescriptor))).allTasks(taskScenarioBuilder -> taskScenarioBuilder.expectZoneId(getZoneFromJobIpAllocation(0, firstIpJobDescriptor))));
String firstJobId = jobsScenarioBuilder.takeJob(0).getJobId();
// Schedule the second task and ensure it's blocked on the first task
jobsScenarioBuilder.schedule(secondIpJobDescriptor, jobScenarioBuilder -> jobScenarioBuilder.template(ScenarioTemplates.jobAccepted()).expectAllTasksCreated().allTasks(taskScenarioBuilder -> taskScenarioBuilder.expectStateUpdates(TaskStatus.TaskState.Accepted)));
String secondJobId = jobsScenarioBuilder.takeJob(1).getJobId();
// Query the gRPC endpoint and ensure the first task does not have a waiting task context field.
TaskQueryResult firstTaskQueryResult = client.findTasks(TaskQuery.newBuilder().setPage(Page.newBuilder().setPageSize(100).build()).putFilteringCriteria("jobIds", firstJobId).build());
assertThat(firstTaskQueryResult.getItemsCount()).isEqualTo(1);
firstTaskQueryResult.getItemsList().forEach(task -> {
assertThat(task.getTaskContextMap()).doesNotContainKeys(TaskAttributes.TASK_ATTRIBUTES_IN_USE_IP_ALLOCATION);
});
String firstTaskId = firstTaskQueryResult.getItems(0).getId();
// Query the gRPC endpoint and ensure the second task has a waiting task context field.
TaskQueryResult secondTaskQueryResult = client.findTasks(TaskQuery.newBuilder().setPage(Page.newBuilder().setPageSize(100).build()).putFilteringCriteria("jobIds", secondJobId).build());
assertThat(secondTaskQueryResult.getItemsCount()).isEqualTo(1);
secondTaskQueryResult.getItemsList().forEach(task -> {
assertThat(task.getTaskContextMap()).contains(new AbstractMap.SimpleImmutableEntry<>(TaskAttributes.TASK_ATTRIBUTES_IN_USE_IP_ALLOCATION, firstTaskId));
});
// Observe the second job and ensure the streamed task has a waiting task context field.
boolean verified = false;
Iterator<JobChangeNotification> it = client.observeJob(JobId.newBuilder().setId(secondJobId).build());
while (it.hasNext()) {
JobChangeNotification jobChangeNotification = it.next();
if (jobChangeNotification.hasTaskUpdate()) {
Map<String, String> taskContext = jobChangeNotification.getTaskUpdate().getTask().getTaskContextMap();
assertThat(taskContext).contains(new AbstractMap.SimpleImmutableEntry<>(TaskAttributes.TASK_ATTRIBUTES_IN_USE_IP_ALLOCATION, firstTaskId));
verified = true;
} else if (jobChangeNotification.hasSnapshotEnd()) {
break;
}
}
assertThat(verified).isTrue();
}
Aggregations