Examples with TestStreamObserver com.netflix.titus.testkit.grpc.TestStreamObserver used on opensource projects

Example 31 with TestStreamObserver

use of com.netflix.titus.testkit.grpc.TestStreamObserver in project titus-control-plane by Netflix.

the class JobScenarioBuilder method updateJobStatus.

public JobScenarioBuilder updateJobStatus(boolean enabled) {
    logger.info("[{}] Changing job {} enable status to {}...", discoverActiveTest(), jobId, enabled);
    Stopwatch stopWatch = Stopwatch.createStarted();
    TestStreamObserver<Empty> responseObserver = new TestStreamObserver<>();
    client.updateJobStatus(JobStatusUpdate.newBuilder().setId(jobId).setEnableStatus(enabled).build(), responseObserver);
    rethrow(() -> responseObserver.awaitDone(TIMEOUT_MS, TimeUnit.MILLISECONDS));
    expectJobUpdateEvent(job -> {
        ServiceJobExt ext = (ServiceJobExt) job.getJobDescriptor().getExtensions();
        return ext.isEnabled() == enabled;
    }, "Job status update did not complete in time");
    logger.info("[{}] Changing job {} enable status to {} finished in {}ms", discoverActiveTest(), jobId, enabled, stopWatch.elapsed(TimeUnit.MILLISECONDS));
    return this;
}

Example 32 with TestStreamObserver

use of com.netflix.titus.testkit.grpc.TestStreamObserver in project titus-control-plane by Netflix.

the class JobObserveTest method observeJobs.

@Test(timeout = LONG_TEST_TIMEOUT_MS)
public void observeJobs() throws InterruptedException {
    TestStreamObserver<JobChangeNotification> subscriber = subscribe(ObserveJobsQuery.getDefaultInstance());
    for (int i = 0; i < 2; i++) {
        String jobId = jobsScenarioBuilder.scheduleAndReturnJob(oneTaskBatchJobDescriptor(), jobScenarioBuilder -> jobScenarioBuilder.template(startTasksInNewJob()).template(ScenarioTemplates.killJob())).getId();
        JobChangeNotification event;
        while (true) {
            event = subscriber.takeNext(SHORT_TIMEOUT_MS, TimeUnit.MILLISECONDS);
            if (event.hasJobUpdate()) {
                Job job = event.getJobUpdate().getJob();
                assertThat(job.getId()).isEqualTo(jobId);
                CellAssertions.assertCellInfo(job, EmbeddedTitusMaster.CELL_NAME);
                if (job.getStatus().getState() == JobState.Finished) {
                    break;
                }
            }
        }
    }
    subscriber.cancel();
}

Example 33 with TestStreamObserver

use of com.netflix.titus.testkit.grpc.TestStreamObserver in project titus-control-plane by Netflix.

the class JobsScenarioBuilder method loadJobs.

private List<JobScenarioBuilder> loadJobs() {
    JobQuery query = JobQuery.newBuilder().setPage(Page.newBuilder().setPageSize(1000)).build();
    Throwable lastFailure = null;
    // since we do not have retry interceptor installed, we make a few attempts directly.
    for (int i = 0; i < 3; i++) {
        try {
            TestStreamObserver<JobQueryResult> responseObserver = new TestStreamObserver<>();
            client.findJobs(query, responseObserver);
            JobQueryResult queryResult = rethrow(() -> responseObserver.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS));
            List<JobScenarioBuilder> result = new ArrayList<>();
            queryResult.getItemsList().forEach(job -> {
                result.add(new JobScenarioBuilder(titusOperations, this, job.getId()));
            });
            return result;
        } catch (Exception e) {
            lastFailure = e;
            logger.info("Cannot load jobs from TitusMaster (might be not ready yet). Waiting 1sec before next try...");
            try {
                Thread.sleep(1_000);
            } catch (InterruptedException ignore) {
            }
        }
    }
    throw new IllegalStateException("Cannot load jobs: " + lastFailure, lastFailure);
}

Example 34 with TestStreamObserver

use of com.netflix.titus.testkit.grpc.TestStreamObserver in project titus-control-plane by Netflix.

the class JobsScenarioBuilder method schedule.

public JobsScenarioBuilder schedule(JobDescriptor jobDescriptor, Function<JobScenarioBuilder, JobScenarioBuilder> jobScenario) {
    TestStreamObserver<JobId> responseObserver = new TestStreamObserver<>();
    client.withDeadlineAfter(TIMEOUT_MS, TimeUnit.MILLISECONDS).createJob(GrpcJobManagementModelConverters.toGrpcJobDescriptor(jobDescriptor), responseObserver);
    JobId jobId;
    try {
        jobId = responseObserver.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    } catch (Exception e) {
        throw new IllegalStateException(e);
    }
    Preconditions.checkNotNull(jobId, "Job create operation not completed in time");
    TestStreamObserver<JobChangeNotification> eventStream = new TestStreamObserver<>();
    client.observeJob(jobId, eventStream);
    JobScenarioBuilder jobScenarioBuilder = new JobScenarioBuilder(titusOperations, this, jobId.getId());
    jobScenarioBuilders.add(jobScenarioBuilder);
    jobScenario.apply(jobScenarioBuilder);
    return this;
}

Example 35 with TestStreamObserver

use of com.netflix.titus.testkit.grpc.TestStreamObserver in project titus-control-plane by Netflix.

the class TaskScenarioBuilder method deleteTaskAttributes.

public TaskScenarioBuilder deleteTaskAttributes(List<String> keys) {
    String taskId = getTask().getId();
    logger.info("[{}] Deleting attributes of task {} of job {}...", discoverActiveTest(), taskId, jobScenarioBuilder.getJobId());
    Stopwatch stopWatch = Stopwatch.createStarted();
    TestStreamObserver<Empty> responseObserver = new TestStreamObserver<>();
    jobClient.deleteTaskAttributes(TaskAttributesDeleteRequest.newBuilder().setTaskId(taskId).addAllKeys(keys).build(), responseObserver);
    rethrow(() -> responseObserver.awaitDone(TIMEOUT_MS, TimeUnit.MILLISECONDS));
    logger.info("[{}] Task {} updated in {}[ms]", discoverActiveTest(), taskId, stopWatch.elapsed(TimeUnit.MILLISECONDS));
    return this;
}