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

Example 36 with TestStreamObserver

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

the class TaskScenarioBuilder method updateTaskAttributes.

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

Example 37 with TestStreamObserver

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

the class TaskScenarioBuilder method moveTask.

public TaskScenarioBuilder moveTask(String targetJobId) {
    String taskId = getTask().getId();
    logger.info("[{}] Moving Task {} to another job {}", discoverActiveTest(), taskId, targetJobId);
    Stopwatch stopWatch = Stopwatch.createStarted();
    TestStreamObserver<Empty> responseObserver = new TestStreamObserver<>();
    jobClient.moveTask(TaskMoveRequest.newBuilder().setSourceJobId(jobScenarioBuilder.getJobId()).setTargetJobId(targetJobId).setTaskId(taskId).build(), responseObserver);
    rethrow(() -> responseObserver.awaitDone(TIMEOUT_MS, TimeUnit.MILLISECONDS));
    logger.info("[{}] Task {} moved to job {} in {}[ms]", discoverActiveTest(), taskId, targetJobId, stopWatch.elapsed(TimeUnit.MILLISECONDS));
    return this;
}

Example 38 with TestStreamObserver

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

the class TaskScenarioBuilder method internalKill.

private TaskScenarioBuilder internalKill(boolean shrink, boolean preventMinSizeUpdate, boolean moveToFinished) {
    String taskId = getTask().getId();
    logger.info("[{}] Killing task: jobId={}, taskId={}, shrink={}, preventMinSizeUpdate={}...", discoverActiveTest(), jobScenarioBuilder.getJobId(), taskId, shrink, preventMinSizeUpdate);
    Stopwatch stopWatch = Stopwatch.createStarted();
    TestStreamObserver<Empty> responseObserver = new TestStreamObserver<>();
    jobClient.killTask(TaskKillRequest.newBuilder().setTaskId(taskId).setShrink(shrink).setPreventMinSizeUpdate(preventMinSizeUpdate).build(), responseObserver);
    rethrow(() -> responseObserver.awaitDone(TIMEOUT_MS, TimeUnit.MILLISECONDS));
    expectStateUpdates(TaskStatus.TaskState.KillInitiated);
    if (moveToFinished) {
        kubeCluster.moveToFinishedSuccess(taskId);
    }
    logger.info("[{}] Task {} killed in {}[ms]", discoverActiveTest(), taskId, stopWatch.elapsed(TimeUnit.MILLISECONDS));
    return this;
}

Example 39 with TestStreamObserver

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

the class AutoScalingGrpcTest method testDeletePolicyById.

/**
 * Test that a policy can be deleted.
 *
 * @throws Exception
 */
@Test(timeout = TEST_TIMEOUT_MS)
public void testDeletePolicyById() throws Exception {
    // Put a policy
    String jobId = "Titus-1";
    PutPolicyRequest putPolicyRequest = AutoScalingTestUtils.generatePutPolicyRequest(jobId, PolicyType.StepScaling);
    TestStreamObserver<ScalingPolicyID> putResponse = new TestStreamObserver<>();
    client.setAutoScalingPolicy(putPolicyRequest, putResponse);
    ScalingPolicyID scalingPolicyID = putResponse.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    assertThat(!scalingPolicyID.getId().isEmpty());
    // Delete the policy
    TestStreamObserver<Empty> deletePolicyResult = new TestStreamObserver<>();
    DeletePolicyRequest deletePolicyRequest = DeletePolicyRequest.newBuilder().setId(scalingPolicyID).build();
    client.deleteAutoScalingPolicy(deletePolicyRequest, deletePolicyResult);
    deletePolicyResult.awaitDone();
    assertThat(deletePolicyResult.hasError()).isFalse();
    // Make sure it's set to Deleting or Deleted state
    TestStreamObserver<GetPolicyResult> getResponse = new TestStreamObserver<>();
    client.getScalingPolicy(scalingPolicyID, getResponse);
    GetPolicyResult getPolicyResult = getResponse.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    log.info("Got result {}", getPolicyResult);
    assertThat(getPolicyResult.getItemsCount()).isEqualTo(1);
    assertThat(isDeletingState(getPolicyResult.getItems(0).getPolicyState()));
}

Example 40 with TestStreamObserver

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

the class AutoScalingGrpcTest method testUpdatePolicyConfigurationForTargetTracking.

/**
 * Test policy configuration update for target tracking policy
 *
 * @throws Exception
 */
@Test(timeout = TEST_TIMEOUT_MS)
public void testUpdatePolicyConfigurationForTargetTracking() throws Exception {
    String jobId = "Titus-123";
    PutPolicyRequest putPolicyRequest = AutoScalingTestUtils.generatePutPolicyRequest(jobId, PolicyType.TargetTrackingScaling);
    TestStreamObserver<ScalingPolicyID> putResponse = new TestStreamObserver<>();
    client.setAutoScalingPolicy(putPolicyRequest, putResponse);
    ScalingPolicyID scalingPolicyID = putResponse.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    assertThat(!scalingPolicyID.getId().isEmpty());
    TestStreamObserver<Empty> updateResponse = new TestStreamObserver<>();
    client.updateAutoScalingPolicy(AutoScalingTestUtils.generateUpdateTargetTrackingPolicyRequest(scalingPolicyID.getId(), 100.0), updateResponse);
    updateResponse.awaitDone();
    AutoScalingPolicyTests.waitForCondition(() -> {
        TestStreamObserver<GetPolicyResult> getResponse = new TestStreamObserver<>();
        client.getScalingPolicy(scalingPolicyID, getResponse);
        try {
            GetPolicyResult getPolicyResult = getResponse.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
            return getPolicyResult.getItemsCount() == 1 && getPolicyResult.getItems(0).getScalingPolicy().getTargetPolicyDescriptor().getTargetValue().getValue() == 100.0;
        } catch (Exception ignored) {
        }
        return false;
    });
    TestStreamObserver<GetPolicyResult> getResponse = new TestStreamObserver<>();
    client.getScalingPolicy(scalingPolicyID, getResponse);
    getResponse.awaitDone();
    GetPolicyResult getPolicyResult = getResponse.takeNext(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    assertThat(getPolicyResult.getItemsCount()).isEqualTo(1);
    DoubleValue targetValue = getPolicyResult.getItems(0).getScalingPolicy().getTargetPolicyDescriptor().getTargetValue();
    assertThat(targetValue.getValue()).isEqualTo(100.0);
}