Java Code Examples for rx.observers.AssertableSubscriber#getOnNextEvents()

@Test
public void getJobLoadBalancersWithOneFailingCell() {
    JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1);
    JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_1, LB_2);
    final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(
            Arrays.asList(jobLoadBalancer1, jobLoadBalancer2));

    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellTwo.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.INTERNAL));

    final AssertableSubscriber<GetJobLoadBalancersResult> resultSubscriber = service.getLoadBalancers(
            JobId.newBuilder().setId(JOB_1).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    final List<GetJobLoadBalancersResult> onNextEvents = resultSubscriber.getOnNextEvents();
    assertThat(onNextEvents).hasSize(1);
    final List<LoadBalancerId> loadBalancersList = onNextEvents.get(0).getLoadBalancersList();
    final List<String> resultLoadBalancers = loadBalancersList.stream().map(loadBalancerId -> loadBalancerId.getId()).collect(Collectors.toList());
    assertThat(resultLoadBalancers).contains(LB_1, LB_2);
}
 

@Test
public void getAllLoadBalancersNoPagination() {
    JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1);
    JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_2, LB_2);
    final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(singletonList(jobLoadBalancer1));
    final CellWithLoadBalancers cellWithLoadBalancersTwo = new CellWithLoadBalancers(singletonList(jobLoadBalancer2));
    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellTwo.getServiceRegistry().addService(cellWithLoadBalancersTwo);

    final AssertableSubscriber<GetAllLoadBalancersResult> resultSubscriber = service.getAllLoadBalancers(
            GetAllLoadBalancersRequest.newBuilder().setPage(Page.newBuilder().setPageNumber(0).setPageSize(10)).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    final List<GetAllLoadBalancersResult> onNextEvents = resultSubscriber.getOnNextEvents();
    assertThat(onNextEvents).hasSize(1);
    assertThat(onNextEvents.get(0).getJobLoadBalancersCount()).isEqualTo(2);
}
 

@Test
public void getPoliciesFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithPolicies cellTwoService = new CellWithPolicies(Collections.singletonList(policyTwoResult));

    cellOne.getServiceRegistry().addService(cellOneService);
    cellTwo.getServiceRegistry().addService(cellTwoService);

    final AssertableSubscriber<GetPolicyResult> testSubscriber = service.getAllScalingPolicies(JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    List<GetPolicyResult> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItemsCount()).isEqualTo(2);
}
 

@Test
public void getLoadBalancersForJob() {
    JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1);
    JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_2, LB_2);
    final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(singletonList(jobLoadBalancer1));
    final CellWithLoadBalancers cellWithLoadBalancersTwo = new CellWithLoadBalancers(singletonList(jobLoadBalancer2));
    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellTwo.getServiceRegistry().addService(cellWithLoadBalancersTwo);

    final AssertableSubscriber<GetJobLoadBalancersResult> resultSubscriber = service.getLoadBalancers(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    resultSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    final List<GetJobLoadBalancersResult> onNextEvents = resultSubscriber.getOnNextEvents();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents.get(0).getLoadBalancersCount()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getLoadBalancers(0).getId()).isEqualTo(LB_2);
}
 

@Test
public void findTasksMergesAllCellsIntoSingleResult() {
    List<Task> cellSnapshot = new ArrayList<>();
    // 10 jobs on each cell with TASKS_IN_GENERATED_JOBS tasks each
    for (int i = 0; i < 5; i++) {
        cellSnapshot.addAll(dataGenerator.newBatchJobWithTasks());
        cellSnapshot.addAll(dataGenerator.newServiceJobWithTasks());
        clock.advanceTime(1, TimeUnit.MINUTES);
    }
    cell.getServiceRegistry().addService(new CellWithFixedTasksService(cellSnapshot));

    TaskQuery query = TaskQuery.newBuilder()
            .setPage(toGrpcPage(Page.unlimited()))
            .build();

    final AssertableSubscriber<TaskQueryResult> testSubscriber = service.findTasks(query, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors().assertCompleted();
    testSubscriber.assertValueCount(1);
    final List<TaskQueryResult> results = testSubscriber.getOnNextEvents();
    assertThat(results).hasSize(1);

    // expect stackName to have changed
    List<Task> expected = cellSnapshot.stream()
            .sorted(JobManagerCursors.taskCursorOrderComparator())
            .map(this::withStackName)
            .collect(Collectors.toList());
    assertThat(results.get(0).getItemsList()).containsExactlyElementsOf(expected);
}
 

@Test
public void findJobsAddsStackName() {
    Random random = new Random();
    final List<Job> cellSnapshot = new ArrayList<>();
    for (int i = 0; i < 5; i++) {
        cellSnapshot.addAll(dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob));
        cellSnapshot.addAll(dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob));
        clock.advanceTime(1, TimeUnit.MINUTES);
    }
    cell.getServiceRegistry().addService(new CellWithFixedJobsService(cellSnapshot, PublishSubject.create()));

    JobQuery query = JobQuery.newBuilder()
            .setPage(toGrpcPage(Page.unlimited()))
            .build();

    final AssertableSubscriber<JobQueryResult> testSubscriber = service.findJobs(query, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors().assertCompleted();
    testSubscriber.assertValueCount(1);
    final List<JobQueryResult> results = testSubscriber.getOnNextEvents();
    assertThat(results).hasSize(1);

    // expect stackName to have changed
    List<Job> expected = cellSnapshot.stream()
            .sorted(JobManagerCursors.jobCursorOrderComparator())
            .map(this::withStackName)
            .collect(Collectors.toList());
    assertThat(results.get(0).getItemsList()).containsExactlyElementsOf(expected);
}
 

@Test
public void findJobsEmptyPage() {
    Random random = new Random();
    final List<Job> cellOneSnapshot = new ArrayList<>();
    final List<Job> cellTwoSnapshot = new ArrayList<>();
    for (int i = 0; i < 5; i++) {
        Iterables.addAll(cellOneSnapshot, Iterables.concat(
                dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob),
                dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        ));
        Iterables.addAll(cellTwoSnapshot, Iterables.concat(
                dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob),
                dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        ));
        clock.advanceTime(1, TimeUnit.MINUTES);
    }
    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(cellOneSnapshot, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(cellTwoSnapshot, cellTwoUpdates.serialize()));

    JobQuery query = JobQuery.newBuilder()
            .setPage(toGrpcPage(Page.empty()))
            .build();

    final AssertableSubscriber<JobQueryResult> testSubscriber = service.findJobs(query, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors().assertCompleted();
    testSubscriber.assertValueCount(1);
    final List<JobQueryResult> results = testSubscriber.getOnNextEvents();
    assertThat(results).hasSize(1);
    assertThat(results.get(0).getItemsList()).isEmpty();
    assertThat(results.get(0).getPagination().getHasMore()).isFalse();
}
 

@Test
public void getPoliciesForJobFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).setJobId(JOB_1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).setJobId(JOB_2).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithPolicies cellTwoService = new CellWithPolicies(Collections.singletonList(policyTwoResult));

    cellOne.getServiceRegistry().addService(cellOneService);
    cellTwo.getServiceRegistry().addService(cellTwoService);

    AssertableSubscriber<GetPolicyResult> testSubscriber = service.getJobScalingPolicies(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();

    testSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItemsCount()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);

    // Bad policy id, currently each Cell returns an empty result
    testSubscriber = service.getJobScalingPolicies(JobId.newBuilder().setId("badID").build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItemsCount()).isEqualTo(0);
}
 

@Test
public void removeLoadBalancer() {
    JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1);
    JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_2, LB_2);
    JobLoadBalancer jobLoadBalancer3 = new JobLoadBalancer(JOB_2, LB_3);
    final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(singletonList(jobLoadBalancer1));
    final CellWithLoadBalancers cellWithLoadBalancersTwo = new CellWithLoadBalancers(new ArrayList<>(Arrays.asList(jobLoadBalancer2, jobLoadBalancer3)));

    final CellWithJobIds cellWithJobIdsOne = new CellWithJobIds(singletonList(JOB_1));
    final CellWithJobIds cellWithJobIdsTwo = new CellWithJobIds(singletonList(JOB_2));

    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellOne.getServiceRegistry().addService(cellWithJobIdsOne);
    cellTwo.getServiceRegistry().addService(cellWithLoadBalancersTwo);
    cellTwo.getServiceRegistry().addService(cellWithJobIdsTwo);

    final AssertableSubscriber<Void> resultSubscriber = service.removeLoadBalancer(
            RemoveLoadBalancerRequest.newBuilder().setJobId(JOB_2).setLoadBalancerId(LoadBalancerId.newBuilder().setId(LB_2).build()).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    resultSubscriber.assertNoValues();
    resultSubscriber.assertCompleted();

    final AssertableSubscriber<GetJobLoadBalancersResult> jobResults = service.getLoadBalancers(
            JobId.newBuilder().setId(JOB_2).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    jobResults.awaitValueCount(1, 1, TimeUnit.SECONDS);
    jobResults.assertNoErrors();
    final List<GetJobLoadBalancersResult> onNextEvents = jobResults.getOnNextEvents();
    assertThat(onNextEvents.size()).isEqualTo(1);
    final List<LoadBalancerId> loadBalancersList = onNextEvents.get(0).getLoadBalancersList();
    assertThat(loadBalancersList.size()).isEqualTo(1);
    assertThat(loadBalancersList.get(0).getId()).isEqualTo(LB_3);
}
 

@Test
public void addLoadBalancer() {
    JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1);
    JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_2, LB_2);
    final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(singletonList(jobLoadBalancer1));
    final CellWithLoadBalancers cellWithLoadBalancersTwo = new CellWithLoadBalancers(new ArrayList<>(singletonList(jobLoadBalancer2)));

    final CellWithJobIds cellWithJobIdsOne = new CellWithJobIds(singletonList(JOB_1));
    final CellWithJobIds cellWithJobIdsTwo = new CellWithJobIds(singletonList(JOB_2));

    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellOne.getServiceRegistry().addService(cellWithJobIdsOne);
    cellTwo.getServiceRegistry().addService(cellWithLoadBalancersTwo);
    cellTwo.getServiceRegistry().addService(cellWithJobIdsTwo);

    final AssertableSubscriber<Void> resultSubscriber = service.addLoadBalancer(
            AddLoadBalancerRequest.newBuilder().setJobId(JOB_2).setLoadBalancerId(LoadBalancerId.newBuilder().setId(LB_3).build()).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    resultSubscriber.assertNoValues();
    resultSubscriber.assertCompleted();

    final AssertableSubscriber<GetJobLoadBalancersResult> jobResults = service.getLoadBalancers(
            JobId.newBuilder().setId(JOB_2).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    jobResults.assertNoErrors();
    final List<GetJobLoadBalancersResult> onNextEvents = jobResults.getOnNextEvents();
    assertThat(onNextEvents.size()).isEqualTo(1);
    final List<LoadBalancerId> loadBalancersList = onNextEvents.get(0).getLoadBalancersList();
    assertThat(loadBalancersList.size()).isEqualTo(2);

    final List<String> resultLoadBalancerIds = loadBalancersList.stream().map(LoadBalancerId::getId).collect(Collectors.toList());
    assertThat(resultLoadBalancerIds).contains(LB_2, LB_3);
}
 

@Test
public void getPolicyByIdFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithPolicies cellTwoService = new CellWithPolicies(Collections.singletonList(policyTwoResult));

    cellOne.getServiceRegistry().addService(cellOneService);
    cellTwo.getServiceRegistry().addService(cellTwoService);

    AssertableSubscriber<GetPolicyResult> testSubscriber = service.getScalingPolicy(
            ScalingPolicyID.newBuilder().setId(POLICY_2).build(), JUNIT_REST_CALL_METADATA).test();

    testSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItemsCount()).isEqualTo(1);

    // Bad id. The current behavior is "INTERNAL: Completed without a response", but it will change to NOT_FOUND someday
    testSubscriber = service.getScalingPolicy(ScalingPolicyID.newBuilder().setId("badID").build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertError(StatusRuntimeException.class);
    testSubscriber.assertNoValues();
    List<Throwable> onErrorEvents = testSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents).isNotNull();
    assertThat(onErrorEvents).hasSize(1);
    assertThat(Status.fromThrowable(onErrorEvents.get(0)).getCode()).isEqualTo(Status.CANCELLED.getCode());
}
 

/**
 * 1. older (bucketized) first
 * 2. in each bucket, larger batches first
 */
@Test
public void emitAccordingToStrategy() {
    final int initialDelay = minimumTimeInQueueMs;
    final Instant now = Instant.ofEpochMilli(testScheduler.now());
    final Instant start = now.minus(ofHours(1));
    final Instant firstBucket = start.plus(ofMillis(timeWindowBucketSizeMs));
    final Instant secondBucket = firstBucket.plus(ofMillis(timeWindowBucketSizeMs));
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(initialDelay);

    final List<Batch<BatchableOperationMock, String>> expected = Arrays.asList(
            // older bucket
            Batch.of("resource1",
                    new BatchableOperationMock(LOW, start, "resource1", "sub1", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(start), "resource1", "sub2", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(firstBucket), "resource1", "sub3", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(secondBucket), "resource1", "sub4", "create")
            ),
            // larger in firstBucket
            Batch.of("resource2",
                    new BatchableOperationMock(LOW, randomWithinBucket(firstBucket), "resource2", "sub1", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(firstBucket), "resource2", "sub2", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(secondBucket), "resource2", "sub3", "create")
            ),
            // larger in secondBucket
            Batch.of("resource3",
                    new BatchableOperationMock(LOW, randomWithinBucket(secondBucket), "resource3", "sub1", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(secondBucket), "resource3", "sub2", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(secondBucket), "resource3", "sub3", "create")
            ),
            // last
            Batch.of("resource4",
                    new BatchableOperationMock(LOW, randomWithinBucket(secondBucket), "resource4", "sub1", "create"),
                    new BatchableOperationMock(LOW, randomWithinBucket(secondBucket), "resource4", "sub2", "create")
            )
    );

    final AssertableSubscriber<Batch<BatchableOperationMock, String>> subscriber = Observable.from(toUpdateList(expected))
            .lift(batcher)
            .test();

    testScheduler.advanceTimeBy(initialDelay, TimeUnit.MILLISECONDS);
    subscriber.assertNoErrors()
            .assertValueCount(expected.size())
            .assertCompleted();

    final List<Batch<BatchableOperationMock, String>> events = subscriber.getOnNextEvents();
    Assertions.assertThat(events).hasSize(expected.size());
    for (int i = 0; i < expected.size(); i++) {
        final Batch<BatchableOperationMock, String> actual = events.get(i);
        final Batch<BatchableOperationMock, String> expectedBatch = expected.get(i);
        final List<BatchableOperationMock> expectedUpdates = expectedBatch.getItems();
        final BatchableOperationMock[] expectedUpdatesArray = expectedUpdates.toArray(new BatchableOperationMock[expectedUpdates.size()]);

        assertThat(actual).isNotNull();
        assertThat(actual.getIndex()).isEqualTo(expectedBatch.getIndex());
        assertThat(actual.getItems()).containsExactlyInAnyOrder(expectedUpdatesArray);
    }
}
 

@Test
public void createAndDeletePolicyForJobIdFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).setJobId(JOB_1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).setJobId(JOB_2).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithPolicies cellTwoService = new CellWithPolicies(Collections.singletonList(policyTwoResult));
    CellWithJobIds cellOneJobsService = new CellWithJobIds(Collections.singletonList(JOB_1));
    CellWithJobIds cellTwoJobsService = new CellWithJobIds(Collections.singletonList(JOB_2));

    cellOne.getServiceRegistry().addService(cellOneService);
    cellTwo.getServiceRegistry().addService(cellOneJobsService);
    cellTwo.getServiceRegistry().addService(cellTwoService);
    cellTwo.getServiceRegistry().addService(cellTwoJobsService);

    AssertableSubscriber<ScalingPolicyID> testSubscriber = service.setAutoScalingPolicy(PutPolicyRequest.newBuilder().setJobId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    assertThat(testSubscriber.getOnErrorEvents().isEmpty()).isTrue();

    List<ScalingPolicyID> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getId()).isNotEmpty();

    String newPolicyId = onNextEvents.get(0).getId();
    AssertableSubscriber<GetPolicyResult> testSubscriber2 = service.getJobScalingPolicies(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents1 = testSubscriber2.getOnNextEvents();
    assertThat(onNextEvents1).isNotNull();
    assertThat(onNextEvents1.size()).isEqualTo(1);
    assertThat(onNextEvents1.get(0).getItemsCount()).isEqualTo(2);
    assertThat(onNextEvents1.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents1.get(0).getItems(1).getJobId()).isEqualTo(JOB_2);

    DeletePolicyRequest deletePolicyRequest = DeletePolicyRequest.newBuilder().setId(ScalingPolicyID.newBuilder().setId(newPolicyId).build()).build();
    AssertableSubscriber<Void> testSubscriber3 = service.deleteAutoScalingPolicy(deletePolicyRequest, JUNIT_REST_CALL_METADATA).test();
    testSubscriber3.assertNoErrors();

    AssertableSubscriber<GetPolicyResult> testSubscriber4 = service.getJobScalingPolicies(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents2 = testSubscriber4.getOnNextEvents();
    assertThat(onNextEvents2).isNotNull();
    assertThat(onNextEvents2.size()).isEqualTo(1);
    assertThat(onNextEvents2.get(0).getItemsCount()).isEqualTo(1);
    assertThat(onNextEvents2.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents2.get(0).getItems(0).getId().getId()).isEqualTo(POLICY_2);
}
 

@Test
public void createAndUpdatePolicyForJobIdFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).setJobId(JOB_1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).setJobId(JOB_2).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithPolicies cellTwoService = new CellWithPolicies(Collections.singletonList(policyTwoResult));

    CellWithJobIds cellOneJobsService = new CellWithJobIds(Collections.singletonList(JOB_1));
    CellWithJobIds cellTwoJobsService = new CellWithJobIds(Collections.singletonList(JOB_2));
    cellOne.getServiceRegistry().addService(cellOneService);
    cellOne.getServiceRegistry().addService(cellOneJobsService);
    cellTwo.getServiceRegistry().addService(cellTwoService);
    cellTwo.getServiceRegistry().addService(cellTwoJobsService);

    AssertableSubscriber<ScalingPolicyID> testSubscriber = service.setAutoScalingPolicy(PutPolicyRequest.newBuilder().setJobId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();

    List<ScalingPolicyID> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getId()).isNotEmpty();

    AssertableSubscriber<GetPolicyResult> testSubscriber2 = service.getJobScalingPolicies(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents1 = testSubscriber2.getOnNextEvents();
    assertThat(onNextEvents1).isNotNull();
    assertThat(onNextEvents1.size()).isEqualTo(1);
    assertThat(onNextEvents1.get(0).getItemsCount()).isEqualTo(2);
    assertThat(onNextEvents1.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents1.get(0).getItems(1).getJobId()).isEqualTo(JOB_2);

    UpdatePolicyRequest updatePolicyRequest = UpdatePolicyRequest.newBuilder()
            .setPolicyId(ScalingPolicyID.newBuilder().setId(POLICY_2))
            .setScalingPolicy(ScalingPolicy.newBuilder()
                    .setTargetPolicyDescriptor(TargetTrackingPolicyDescriptor.newBuilder()
                            .setTargetValue(DoubleValue.newBuilder().setValue(100.0).build()).build()).build()).build();

    AssertableSubscriber<Void> testSubscriber3 = service.updateAutoScalingPolicy(updatePolicyRequest, JUNIT_REST_CALL_METADATA).test();
    testSubscriber3.assertNoErrors();

    AssertableSubscriber<GetPolicyResult> testSubscriber4 = service.getScalingPolicy(ScalingPolicyID.newBuilder().setId(POLICY_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents2 = testSubscriber4.getOnNextEvents();
    assertThat(onNextEvents2).isNotNull();
    assertThat(onNextEvents2.size()).isEqualTo(1);
    assertThat(onNextEvents2.get(0).getItemsCount()).isEqualTo(1);
    assertThat(onNextEvents2.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents2.get(0).getItems(0).getId().getId()).isEqualTo(POLICY_2);
    ScalingPolicy scalingPolicy = onNextEvents2.get(0).getItems(0).getScalingPolicy();
    double updatedValue = scalingPolicy.getTargetPolicyDescriptor().getTargetValue().getValue();
    assertThat(updatedValue).isEqualTo(100);
}
 

@Test
public void observeJobsWaitsForAllMarkers() {
    final List<Job> cellOneSnapshot = Arrays.asList(
            Job.newBuilder().setId("cell-1-job-1").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-1-job-2").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-1-job-3").setStatus(ACCEPTED_STATE).build()
    );
    final List<Job> cellTwoSnapshot = Arrays.asList(
            Job.newBuilder().setId("cell-2-job-1").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-2-job-2").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-2-job-3").setStatus(ACCEPTED_STATE).build()
    );

    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(cellOneSnapshot, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(cellTwoSnapshot, cellTwoUpdates.serialize()));

    final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test();
    List<JobChangeNotification> expected = Stream.concat(
            cellOneSnapshot.stream().map(this::toNotification).map(this::withStackName),
            cellTwoSnapshot.stream().map(this::toNotification).map(this::withStackName)
    ).collect(Collectors.toList());
    // single marker for all cells
    final JobChangeNotification mergedMarker = JobChangeNotification.newBuilder().setSnapshotEnd(SnapshotEnd.newBuilder()).build();
    expected.add(mergedMarker);

    testSubscriber.awaitValueCount(7, 1, TimeUnit.SECONDS);
    List<JobChangeNotification> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).last().isEqualTo(mergedMarker);
    assertThat(onNextEvents).containsExactlyInAnyOrder(expected.toArray(new JobChangeNotification[expected.size()]));

    // more updates are flowing
    final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-10").setStatus(ACCEPTED_STATE).build());
    final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-10").setStatus(ACCEPTED_STATE).build());
    cellOneUpdates.onNext(cellOneUpdate);
    cellTwoUpdates.onNext(cellTwoUpdate);

    testSubscriber.awaitValueCount(9, 1, TimeUnit.SECONDS);
    onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).last().isNotEqualTo(mergedMarker);
    assertThat(onNextEvents).contains(withStackName(cellOneUpdate), withStackName(cellTwoUpdate));
}
 

@Test
public void findJobsMergesAllCellsIntoSingleResult() {
    Random random = new Random();
    final List<Job> cellOneSnapshot = new ArrayList<>();
    final List<Job> cellTwoSnapshot = new ArrayList<>();
    for (int i = 0; i < 5; i++) {
        cellOneSnapshot.addAll(
                dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        cellTwoSnapshot.addAll(
                dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        cellOneSnapshot.addAll(
                dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        cellTwoSnapshot.addAll(
                dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        clock.advanceTime(1, TimeUnit.MINUTES);
    }

    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(cellOneSnapshot, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(cellTwoSnapshot, cellTwoUpdates.serialize()));

    JobQuery query = JobQuery.newBuilder()
            .setPage(toGrpcPage(Page.unlimited()))
            .build();

    final AssertableSubscriber<JobQueryResult> testSubscriber = service.findJobs(query, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors().assertCompleted();
    testSubscriber.assertValueCount(1);
    final List<JobQueryResult> results = testSubscriber.getOnNextEvents();
    assertThat(results).hasSize(1);

    // expect stackName to have changed
    List<Job> expected = Stream.concat(cellOneSnapshot.stream(), cellTwoSnapshot.stream())
            .sorted(JobManagerCursors.jobCursorOrderComparator())
            .map(this::withStackName)
            .collect(Collectors.toList());
    assertThat(results.get(0).getItemsList()).containsExactlyElementsOf(expected);
}
 

/**
 * @param <Q> query type
 * @param <R> result type
 * @param <T> type of items in the result
 */
static <Q, R, T> List<T> walkAllPages(int pageWalkSize,
                                      Function<Q, Observable<R>> pageFetcher,
                                      Function<com.netflix.titus.grpc.protogen.Page, Q> queryFactory,
                                      Function<R, com.netflix.titus.grpc.protogen.Pagination> paginationGetter,
                                      Function<R, List<T>> itemsGetter) {
    List<T> allItems = new ArrayList<>();
    Optional<R> lastResult = Optional.empty();
    int currentCursorPosition = -1;
    int currentPageNumber = 0;

    while (lastResult.map(r -> paginationGetter.apply(r).getHasMore()).orElse(true)) {
        com.netflix.titus.grpc.protogen.Page.Builder builder = com.netflix.titus.grpc.protogen.Page.newBuilder().setPageSize(pageWalkSize);
        if (lastResult.isPresent()) {
            builder.setCursor(paginationGetter.apply(lastResult.get()).getCursor());
        } else {
            builder.setPageNumber(0);
        }

        Q query = queryFactory.apply(builder.build());
        AssertableSubscriber<R> testSubscriber = pageFetcher.apply(query).test();
        testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
        testSubscriber.assertNoErrors().assertCompleted();
        testSubscriber.assertValueCount(1);

        final List<R> results = testSubscriber.getOnNextEvents();
        assertThat(results).hasSize(1);
        R result = results.get(0);
        List<T> items = itemsGetter.apply(result);
        com.netflix.titus.grpc.protogen.Pagination pagination = paginationGetter.apply(result);
        if (pagination.getHasMore()) {
            assertThat(items).hasSize(pageWalkSize);
        }
        currentCursorPosition += items.size();
        if (pagination.getTotalItems() > 0) {
            assertThat(pagination.getCursorPosition()).isEqualTo(currentCursorPosition);
        } else {
            assertThat(pagination.getCursorPosition()).isEqualTo(0);
        }
        assertThat(pagination.getCurrentPage().getPageNumber()).isEqualTo(currentPageNumber++);
        allItems.addAll(items);
        lastResult = Optional.of(result);
        testSubscriber.unsubscribe();
    }

    return allItems;
}