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

Example 41 with TestStreamObserver

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

the class TaskRelocationIntegrationTest method findRelocationPlan.

private Optional<TaskRelocationPlan> findRelocationPlan(String taskId) {
    TaskRelocationPlans plans;
    try {
        TestStreamObserver<TaskRelocationPlans> events = new TestStreamObserver<>();
        client.getCurrentTaskRelocationPlans(TaskRelocationQuery.getDefaultInstance(), events);
        plans = events.getLast();
    } catch (Exception e) {
        if (e.getMessage().contains("Relocation workflow not ready yet")) {
            return Optional.empty();
        }
        throw new RuntimeException(e);
    }
    Optional<TaskRelocationPlan> taskPlan = plans.getPlansList().stream().filter(p -> p.getTaskId().equals(taskId)).findFirst();
    if (taskPlan.isPresent()) {
        return taskPlan;
    }
    // Check if already processed
    try {
        TestStreamObserver<TaskRelocationExecution> events = new TestStreamObserver<>();
        client.getTaskRelocationResult(RelocationTaskId.newBuilder().setId(taskId).build(), events);
        return Optional.of(events.getLast().getTaskRelocationPlan());
    } catch (Exception e) {
        Status status = Status.fromThrowable(e);
        if (status.getCode() == Status.Code.NOT_FOUND) {
            return Optional.empty();
        }
        throw new RuntimeException(e);
    }
}

Example 42 with TestStreamObserver

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

the class LoadBalancerTests method getAllLoadBalancers.

/**
 * Common testing helper that gets all load balancers for a given page range, ensures the gRPC
 * request was successful, and returns the page result.
 */
public static GetAllLoadBalancersResult getAllLoadBalancers(Supplier<Page> pageSupplier, BiConsumer<GetAllLoadBalancersRequest, TestStreamObserver<GetAllLoadBalancersResult>> getAllLoadBalancers) {
    GetAllLoadBalancersRequest request = GetAllLoadBalancersRequest.newBuilder().setPage(pageSupplier.get()).build();
    TestStreamObserver<GetAllLoadBalancersResult> getResponse = new TestStreamObserver<>();
    getAllLoadBalancers.accept(request, getResponse);
    GetAllLoadBalancersResult result = null;
    try {
        result = getResponse.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    } catch (Exception e) {
        logger.error("Exception in getAllLoadBalancers", e);
        assert false;
    }
    assertThat(getResponse.hasError()).isFalse();
    return result;
}