Search in sources :

Example 1 with TaskScenarioBuilder

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();
}
Also used : TaskScenarioBuilder(com.netflix.titus.master.integration.v3.scenario.TaskScenarioBuilder) JobDescriptorGenerator.serviceJobDescriptors(com.netflix.titus.testkit.model.job.JobDescriptorGenerator.serviceJobDescriptors) Assertions.assertThat(org.assertj.core.api.Assertions.assertThat) ScenarioTemplates(com.netflix.titus.master.integration.v3.scenario.ScenarioTemplates) Function(java.util.function.Function) TaskStatus(com.netflix.titus.grpc.protogen.TaskStatus) SignedIpAddressAllocation(com.netflix.titus.api.jobmanager.model.job.vpc.SignedIpAddressAllocation) Map(java.util.Map) BatchJobExt(com.netflix.titus.api.jobmanager.model.job.ext.BatchJobExt) GrpcJobManagementModelConverters.toGrpcJobDescriptor(com.netflix.titus.runtime.endpoint.v3.grpc.GrpcJobManagementModelConverters.toGrpcJobDescriptor) JobTestUtils.submitBadJob(com.netflix.titus.master.integration.v3.job.JobTestUtils.submitBadJob) TaskQuery(com.netflix.titus.grpc.protogen.TaskQuery) Before(org.junit.Before) BaseIntegrationTest(com.netflix.titus.master.integration.BaseIntegrationTest) JobDescriptor(com.netflix.titus.api.jobmanager.model.job.JobDescriptor) Iterator(java.util.Iterator) ServiceJobExt(com.netflix.titus.api.jobmanager.model.job.ext.ServiceJobExt) TaskQueryResult(com.netflix.titus.grpc.protogen.TaskQueryResult) JobFunctions(com.netflix.titus.api.jobmanager.model.job.JobFunctions) JobManagementServiceGrpc(com.netflix.titus.grpc.protogen.JobManagementServiceGrpc) IntegrationTest(com.netflix.titus.testkit.junit.category.IntegrationTest) Test(org.junit.Test) Category(org.junit.experimental.categories.Category) Page(com.netflix.titus.grpc.protogen.Page) RuleChain(org.junit.rules.RuleChain) TitusStackResource(com.netflix.titus.testkit.junit.master.TitusStackResource) AbstractMap(java.util.AbstractMap) List(java.util.List) Rule(org.junit.Rule) JobDescriptorGenerator.batchJobDescriptors(com.netflix.titus.testkit.model.job.JobDescriptorGenerator.batchJobDescriptors) Ignore(org.junit.Ignore) JobIpAllocationGenerator(com.netflix.titus.testkit.model.job.JobIpAllocationGenerator) EmbeddedTitusCells.basicKubeCell(com.netflix.titus.testkit.embedded.cell.EmbeddedTitusCells.basicKubeCell) TaskAttributes(com.netflix.titus.api.jobmanager.TaskAttributes) JobsScenarioBuilder(com.netflix.titus.master.integration.v3.scenario.JobsScenarioBuilder) JobChangeNotification(com.netflix.titus.grpc.protogen.JobChangeNotification) JobId(com.netflix.titus.grpc.protogen.JobId) AbstractMap(java.util.AbstractMap) JobChangeNotification(com.netflix.titus.grpc.protogen.JobChangeNotification) ServiceJobExt(com.netflix.titus.api.jobmanager.model.job.ext.ServiceJobExt) TaskQueryResult(com.netflix.titus.grpc.protogen.TaskQueryResult) Ignore(org.junit.Ignore) BaseIntegrationTest(com.netflix.titus.master.integration.BaseIntegrationTest) IntegrationTest(com.netflix.titus.testkit.junit.category.IntegrationTest) Test(org.junit.Test)

Aggregations

TaskAttributes (com.netflix.titus.api.jobmanager.TaskAttributes)1 JobDescriptor (com.netflix.titus.api.jobmanager.model.job.JobDescriptor)1 JobFunctions (com.netflix.titus.api.jobmanager.model.job.JobFunctions)1 BatchJobExt (com.netflix.titus.api.jobmanager.model.job.ext.BatchJobExt)1 ServiceJobExt (com.netflix.titus.api.jobmanager.model.job.ext.ServiceJobExt)1 SignedIpAddressAllocation (com.netflix.titus.api.jobmanager.model.job.vpc.SignedIpAddressAllocation)1 JobChangeNotification (com.netflix.titus.grpc.protogen.JobChangeNotification)1 JobId (com.netflix.titus.grpc.protogen.JobId)1 JobManagementServiceGrpc (com.netflix.titus.grpc.protogen.JobManagementServiceGrpc)1 Page (com.netflix.titus.grpc.protogen.Page)1 TaskQuery (com.netflix.titus.grpc.protogen.TaskQuery)1 TaskQueryResult (com.netflix.titus.grpc.protogen.TaskQueryResult)1 TaskStatus (com.netflix.titus.grpc.protogen.TaskStatus)1 BaseIntegrationTest (com.netflix.titus.master.integration.BaseIntegrationTest)1 JobTestUtils.submitBadJob (com.netflix.titus.master.integration.v3.job.JobTestUtils.submitBadJob)1 JobsScenarioBuilder (com.netflix.titus.master.integration.v3.scenario.JobsScenarioBuilder)1 ScenarioTemplates (com.netflix.titus.master.integration.v3.scenario.ScenarioTemplates)1 TaskScenarioBuilder (com.netflix.titus.master.integration.v3.scenario.TaskScenarioBuilder)1 GrpcJobManagementModelConverters.toGrpcJobDescriptor (com.netflix.titus.runtime.endpoint.v3.grpc.GrpcJobManagementModelConverters.toGrpcJobDescriptor)1 EmbeddedTitusCells.basicKubeCell (com.netflix.titus.testkit.embedded.cell.EmbeddedTitusCells.basicKubeCell)1