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

Example 16 with TestStreamObserver

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

the class JobScenarioBuilder method findTasks.

public JobScenarioBuilder findTasks(TaskQuery taskQuery, Predicate<List<com.netflix.titus.grpc.protogen.Task>> tasksPredicate) {
    TestStreamObserver<TaskQueryResult> responseObserver = new TestStreamObserver<>();
    client.findTasks(taskQuery, responseObserver);
    TaskQueryResult result = rethrow(() -> responseObserver.getLast(TIMEOUT_MS, TimeUnit.MILLISECONDS));
    if (result == null) {
        throw new IllegalStateException("TaskQueryResult is null");
    } else if (!tasksPredicate.test(result.getItemsList())) {
        throw new IllegalStateException("Tasks predicate is false");
    }
    return this;
}

Example 17 with TestStreamObserver

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

the class LoadBalancerTests method putLoadBalancersPerJob.

/**
 * Common testing helper that adds a specified number of random load balancer ids to
 * a specified number of jobs. The helper ensures the gRPC request was successful and
 * returns the job ids and load balancer ids as a map.
 */
public static Map<String, Set<LoadBalancerId>> putLoadBalancersPerJob(int numJobs, int numLoadBalancersPerJob, BiConsumer<AddLoadBalancerRequest, TestStreamObserver<Empty>> putLoadBalancer) {
    // Create job entries
    Map<String, Set<LoadBalancerId>> jobIdToLoadBalancersMap = new ConcurrentHashMap<>();
    for (int i = 1; i <= numJobs; i++) {
        jobIdToLoadBalancersMap.put("Titus-" + i, new HashSet<>());
    }
    // For each job, insert load balancers
    jobIdToLoadBalancersMap.forEach((jobId, loadBalancerSet) -> {
        for (int i = 0; i < numLoadBalancersPerJob; i++) {
            LoadBalancerId loadBalancerId = LoadBalancerId.newBuilder().setId(RandomStringUtils.randomAlphanumeric(10)).build();
            AddLoadBalancerRequest request = AddLoadBalancerRequest.newBuilder().setJobId(jobId).setLoadBalancerId(loadBalancerId).build();
            TestStreamObserver<Empty> addResponse = new TestStreamObserver<>();
            putLoadBalancer.accept(request, addResponse);
            assertThatCode(addResponse::awaitDone).doesNotThrowAnyException();
            assertThat(addResponse.hasError()).isFalse();
            loadBalancerSet.add(loadBalancerId);
        }
    });
    return jobIdToLoadBalancersMap;
}

Example 18 with TestStreamObserver

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

the class LoadBalancerTests method removeLoadBalancerFromJob.

/**
 * Common testing helper that removes a load balancer id from a job. The helper ensures the
 * gRPC request was successful.
 */
public static void removeLoadBalancerFromJob(String jobId, LoadBalancerId loadBalancerId, BiConsumer<RemoveLoadBalancerRequest, TestStreamObserver<Empty>> removeLoadBalancers) {
    RemoveLoadBalancerRequest request = RemoveLoadBalancerRequest.newBuilder().setJobId(jobId).setLoadBalancerId(loadBalancerId).build();
    TestStreamObserver<Empty> removeResponse = new TestStreamObserver<>();
    removeLoadBalancers.accept(request, removeResponse);
    assertThatCode(removeResponse::awaitDone).doesNotThrowAnyException();
    assertThat(removeResponse.hasError()).isFalse();
}

Example 19 with TestStreamObserver

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

the class LoadBalancerTests method getLoadBalancersForJob.

/**
 * Common testing helper that gets load balancers for a job, ensures the gRPC request was
 * successful, and returns the load balancer ids as a set.
 */
public static Set<LoadBalancerId> getLoadBalancersForJob(String jobIdStr, BiConsumer<JobId, TestStreamObserver<GetJobLoadBalancersResult>> getJobLoadBalancers) {
    JobId jobId = JobId.newBuilder().setId(jobIdStr).build();
    TestStreamObserver<GetJobLoadBalancersResult> getResponse = new TestStreamObserver<>();
    getJobLoadBalancers.accept(jobId, getResponse);
    GetJobLoadBalancersResult result = null;
    try {
        result = getResponse.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    } catch (Exception e) {
        logger.error("Exception in getLoadBalancersForJob", e);
        assert false;
    }
    assertThat(getResponse.hasError()).isFalse();
    return new HashSet<>(result.getLoadBalancersList());
}

Example 20 with TestStreamObserver

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

the class TaskRelocationIntegrationTest method testEvents.

@Test(timeout = 60_000)
public void testEvents() throws InterruptedException {
    TestStreamObserver<RelocationEvent> events = new TestStreamObserver<>();
    client.observeRelocationEvents(TaskRelocationQuery.getDefaultInstance(), events);
    RelocationEvent firstEvent = events.takeNext(30, TimeUnit.SECONDS);
    assertThat(firstEvent).isNotNull();
    assertThat(firstEvent.getEventCase()).isEqualTo(EventCase.SNAPSHOTEND);
    Task task = createAndPlaceOneTaskJob(TestDataFactory.REMOVABLE_INSTANCE_GROUP_ID);
    relocationConnectorStubs.setQuota(task.getJobId(), 1);
    RelocationEvent secondEvent = events.takeNext(30, TimeUnit.SECONDS);
    assertThat(secondEvent).isNotNull();
    assertThat(secondEvent.getEventCase()).isEqualTo(EventCase.TASKRELOCATIONPLANUPDATEEVENT);
    RelocationEvent thirdEvent = events.takeNext(30, TimeUnit.SECONDS);
    assertThat(thirdEvent).isNotNull();
    assertThat(thirdEvent.getEventCase()).isEqualTo(EventCase.TASKRELOCATIONPLANREMOVEEVENT);
}