rx.observers.AssertableSubscriber Java Examples

@Test
public void observeJobsStopsWhenAnyClientsTerminate() {
    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize()));

    final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test();

    final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-100").setStatus(ACCEPTED_STATE).build());
    final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-200").setStatus(ACCEPTED_STATE).build());
    cellOneUpdates.onNext(cellOneUpdate);
    cellTwoUpdates.onNext(cellTwoUpdate);

    testSubscriber.awaitValueCount(2, 1, TimeUnit.SECONDS);
    assertThat(testSubscriber.getOnErrorEvents()).isEmpty();
    assertThat(testSubscriber.isUnsubscribed()).isFalse();
    assertThat(testSubscriber.getCompletions()).isEqualTo(0);

    // a client finishes
    cellTwoUpdates.onCompleted();

    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    assertThat(testSubscriber.getOnErrorEvents()).isEmpty();
    assertThat(testSubscriber.isUnsubscribed()).isTrue();
    assertThat(testSubscriber.getCompletions()).isEqualTo(1);
}
 

@Test
public void pendingItemsAreFlushedAfterUpstreamCompletes() {
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(minimumTimeInQueueMs);

    final Instant now = Instant.ofEpochMilli(testScheduler.now());
    final BatchableOperationMock first = new BatchableOperationMock(LOW, now.minus(ofSeconds(5)), "resource1", "sub1", "create");
    final BatchableOperationMock second = new BatchableOperationMock(LOW, now.minus(ofSeconds(3)), "resource2", "sub1", "create");
    final Observable<BatchableOperationMock> updates = Observable.from(Arrays.asList(first, second));

    final AssertableSubscriber<Batch<BatchableOperationMock, String>> subscriber = updates.lift(batcher).test();
    testScheduler.advanceTimeBy(1, TimeUnit.MINUTES);
    //noinspection unchecked
    subscriber.assertNoErrors()
            .assertValueCount(2)
            .assertValues(
                    Batch.of("resource1", Collections.singletonList(first)),
                    Batch.of("resource2", Collections.singletonList(second))
            )
            .assertCompleted();
}
 

@Test
public void onCompletedIsNotSentAfterOnError() {
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(minimumTimeInQueueMs);
    final Subject<BatchableOperationMock, BatchableOperationMock> updates = PublishSubject.<BatchableOperationMock>create().toSerialized();

    final AssertableSubscriber<Batch<BatchableOperationMock, String>> subscriber = updates.lift(batcher).test();
    testScheduler.triggerActions();
    subscriber.assertNoTerminalEvent().assertNoValues();

    updates.onError(new RuntimeException("some problem"));
    testScheduler.triggerActions(); // onError is forwarded right away (i.e.: don't wait for the next flush event)
    subscriber.assertNotCompleted().assertError(RuntimeException.class);

    updates.onCompleted();
    subscriber.assertNotCompleted();
}
 

@Test
public void asyncActionSingle() throws Exception {
    AmazonWebServiceRequest someRequest = AmazonWebServiceRequest.NOOP;
    final MockAsyncClient<AmazonWebServiceRequest, String> client = new MockAsyncClient<>(someRequest, "some response");
    Single<String> single = AwsObservableExt.asyncActionSingle(supplier -> client.someAsyncOperation(supplier.handler()));

    TestScheduler testScheduler = Schedulers.test();
    final AssertableSubscriber<String> subscriber = single.subscribeOn(testScheduler).test();

    testScheduler.triggerActions();
    subscriber.assertNoValues();
    subscriber.assertNotCompleted();

    client.run();
    testScheduler.triggerActions();
    subscriber.assertValueCount(1);
    subscriber.assertValue("some response");
    subscriber.assertCompleted();
}
 

@Test
public void minimumTimeInQueue() {
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(1);

    final Instant now = Instant.ofEpochMilli(testScheduler.now());
    final List<BatchableOperationMock> updatesList = Arrays.asList(
            new BatchableOperationMock(LOW, now, "resource1", "sub2", "create"),
            new BatchableOperationMock(LOW, now.minus(ofMillis(minimumTimeInQueueMs + 1)), "resource1", "sub1", "create")
    );

    final AssertableSubscriber<Batch<BatchableOperationMock, String>> subscriber = Observable.from(updatesList).lift(batcher).test();
    testScheduler.triggerActions();
    subscriber.assertNoTerminalEvent().assertNoValues();

    testScheduler.advanceTimeBy(1, TimeUnit.MILLISECONDS);
    subscriber.assertNoErrors()
            .assertValueCount(1)
            .assertValue(Batch.of("resource1", updatesList))
            .assertCompleted();
}
 

@Test
public void onErrorIsNotSentAfterOnCompleted() {
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(minimumTimeInQueueMs);
    final Subject<BatchableOperationMock, BatchableOperationMock> updates = PublishSubject.<BatchableOperationMock>create().toSerialized();

    final AssertableSubscriber<Batch<BatchableOperationMock, String>> subscriber = updates.lift(batcher).test();
    testScheduler.triggerActions();
    subscriber.assertNoTerminalEvent().assertNoValues();

    updates.onCompleted();
    // onCompleted is sent after pending items are drained
    testScheduler.advanceTimeBy(minimumTimeInQueueMs, TimeUnit.MILLISECONDS);
    subscriber.assertNoErrors().assertCompleted();

    updates.onError(new RuntimeException("some problem"));
    subscriber.assertNoErrors();
}
 

@Test
public void shouldSupportStartingALoopWithAnInit() throws Exception {
  MobiusLoop.Builder<String, Integer, Boolean> withInit =
      builder.init(
          new Init<String, Boolean>() {
            @Nonnull
            @Override
            public First<String, Boolean> init(String model) {
              return First.first(model + "-init");
            }
          });

  Observable.Transformer<Integer, String> transformer = RxMobius.loopFrom(withInit, "hi");

  final AssertableSubscriber<String> observer = Observable.just(10).compose(transformer).test();

  waitForSubscriberValueCount(observer, 2);
  observer.assertValues("hi-init", "hi-init10");
}
 

@Test
public void emptyBatchesAreFilteredOut() {
    defaultStubs();

    LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS);
    DefaultLoadBalancerService service = new DefaultLoadBalancerService(
            runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler);

    AssertableSubscriber<Batch<TargetStateBatchable, String>> testSubscriber = service.events().test();

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();

    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();
}
 

@Test
public void testBootstrapEngineInitialization() {
    InternalReconciliationEngine<SimpleReconcilerEvent> bootstrapEngine = mock(InternalReconciliationEngine.class);
    PublishSubject<SimpleReconcilerEvent> eventSubject = PublishSubject.create();
    when(bootstrapEngine.events()).thenReturn(eventSubject);
    when(bootstrapEngine.triggerActions()).thenReturn(true);
    when(bootstrapEngine.getReferenceView()).thenReturn(EntityHolder.newRoot("myRoot1", "myEntity1"));

    DefaultReconciliationFramework<SimpleReconcilerEvent> framework = new DefaultReconciliationFramework<>(
            Collections.singletonList(bootstrapEngine),
            engineFactory,
            IDLE_TIMEOUT_MS,
            ACTIVE_TIMEOUT_MS,
            indexComparators,
            new DefaultRegistry(),
            Optional.of(testScheduler)
    );
    framework.start();
    AssertableSubscriber<SimpleReconcilerEvent> eventSubscriber = framework.events().test();

    eventSubject.onNext(new SimpleReconcilerEvent(EventType.Changed, "test", Optional.empty()));
    eventSubscriber.assertValueCount(1);
}
 

@Test
public void asyncActionCompletable() throws Exception {
    AmazonWebServiceRequest someRequest = AmazonWebServiceRequest.NOOP;
    final MockAsyncClient<AmazonWebServiceRequest, String> client = new MockAsyncClient<>(someRequest, "some response");
    final Completable completable = AwsObservableExt.asyncActionCompletable(factory -> client.someAsyncOperation(factory.handler(
            (req, resp) -> {
                assertEquals(someRequest, req);
                assertEquals("some response", resp);
            },
            (t) -> {
                throw new IllegalStateException("Should never be here");
            }
    )));

    TestScheduler testScheduler = Schedulers.test();
    final AssertableSubscriber<Void> subscriber = completable.subscribeOn(testScheduler).test();

    testScheduler.triggerActions();
    subscriber.assertNotCompleted();

    client.run();
    testScheduler.triggerActions();
    subscriber.assertCompleted();
}
 

private void assertEmitSingleAfterReceiving(BatchableOperationMock expected, BatchableOperationMock... receiving) {
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(minimumTimeInQueueMs);
    final Subject<BatchableOperationMock, BatchableOperationMock> updates = PublishSubject.<BatchableOperationMock>create().toSerialized();

    final AssertableSubscriber<Batch<BatchableOperationMock, String>> subscriber = updates.lift(batcher).test();
    testScheduler.triggerActions();
    subscriber.assertNoTerminalEvent().assertNoValues();

    for (BatchableOperationMock received : receiving) {
        updates.onNext(received);
    }

    testScheduler.advanceTimeBy(2 * minimumTimeInQueueMs, TimeUnit.MILLISECONDS);
    subscriber.assertNoErrors()
            .assertValueCount(1)
            .assertValue(Batch.of(expected.getResourceId(), Collections.singletonList(expected)));
}
 

@Test
public void ignoreErrorsFromDownstream() {
    final Instant now = Instant.ofEpochMilli(testScheduler.now());
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(minimumTimeInQueueMs);
    final Subject<BatchableOperationMock, BatchableOperationMock> updates = PublishSubject.<BatchableOperationMock>create().toSerialized();

    final AssertableSubscriber<?> subscriber = updates.lift(batcher)
            .lift(new ExceptionThrowingOperator("some error happened"))
            .test();
    testScheduler.triggerActions();
    subscriber.assertNoTerminalEvent().assertNoValues();

    for (int i = 0; i < 10; i++) {
        updates.onNext(new BatchableOperationMock(LOW, now, "resource2", "sub2", "create"));
        testScheduler.advanceTimeBy(minimumTimeInQueueMs, TimeUnit.MILLISECONDS);
        subscriber.assertNoTerminalEvent().assertNoValues();
    }

    updates.onCompleted();
    testScheduler.advanceTimeBy(minimumTimeInQueueMs, TimeUnit.MILLISECONDS);
    subscriber.assertNoValues().assertCompleted();
}
 

@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 getLoadBalancersWithOneFailingCell() {
    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<GetAllLoadBalancersResult> resultSubscriber = service.getAllLoadBalancers(
            GetAllLoadBalancersRequest.newBuilder().setPage(Page.newBuilder().setPageSize(10)).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoValues();
    final List<Throwable> onErrorEvents = resultSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents).hasSize(1);
    assertThat(Status.fromThrowable(onErrorEvents.get(0))).isEqualTo(Status.INTERNAL);
}
 

@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 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 allCellsFailing() {
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(failing("one")));
    cellTwo.getServiceRegistry().addService(new CellWithHealthStatus(failing("two")));

    AssertableSubscriber<HealthCheckResponse> subscriber = service.check(HealthCheckRequest.newBuilder().build()).test();
    subscriber.awaitTerminalEvent(10, TimeUnit.SECONDS);
    subscriber.assertNoErrors();
    subscriber.assertValueCount(1);
    HealthCheckResponse response = subscriber.getOnNextEvents().get(0);
    assertThat(response.getStatus()).isEqualTo(NOT_SERVING);
    assertThat(response.getDetailsCount()).isEqualTo(2);
    assertThat(response.getDetails(0).hasDetails()).isTrue();
    assertThat(response.getDetails(1).hasDetails()).isTrue();
    Set<String> cellsSeen = response.getDetailsList().stream()
            .map(s -> s.getDetails().getCell())
            .collect(Collectors.toSet());
    assertThat(cellsSeen).contains("one", "two");
}
 

@Test
public void grpcErrors() {
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(ok("one")));
    cellTwo.getServiceRegistry().addService(new HealthGrpc.HealthImplBase() {
        @Override
        public void check(HealthCheckRequest request, StreamObserver<HealthCheckResponse> responseObserver) {
            responseObserver.onError(Status.DEADLINE_EXCEEDED.asRuntimeException());
        }
    });

    AssertableSubscriber<HealthCheckResponse> subscriber = service.check(HealthCheckRequest.newBuilder().build()).test();
    subscriber.awaitTerminalEvent(10, TimeUnit.SECONDS);
    subscriber.assertNoErrors();
    subscriber.assertValueCount(1);
    HealthCheckResponse response = subscriber.getOnNextEvents().get(0);
    assertThat(response.getStatus()).isEqualTo(NOT_SERVING);
    assertThat(response.getDetailsCount()).isEqualTo(2);
    List<ServerStatus> errors = response.getDetailsList().stream()
            .filter(ServerStatus::hasError)
            .collect(Collectors.toList());
    assertThat(errors).hasSize(1);
    assertThat(errors.get(0).getError().getCell()).isEqualTo("two");
    assertThat(errors.get(0).getError().getErrorCode()).isEqualTo(DEADLINE_EXCEEDED.toString());
}
 

@Test
public void oneCellFailing() {
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(ok("one")));
    cellTwo.getServiceRegistry().addService(new CellWithHealthStatus(failing("two")));

    AssertableSubscriber<HealthCheckResponse> subscriber = service.check(HealthCheckRequest.newBuilder().build()).test();
    subscriber.awaitTerminalEvent(10, TimeUnit.SECONDS);
    subscriber.assertNoErrors();
    subscriber.assertValueCount(1);
    HealthCheckResponse response = subscriber.getOnNextEvents().get(0);
    assertThat(response.getStatus()).isEqualTo(NOT_SERVING);
    assertThat(response.getDetailsCount()).isEqualTo(2);
    assertThat(response.getDetails(0).hasDetails()).isTrue();
    assertThat(response.getDetails(1).hasDetails()).isTrue();
    Set<String> cellsSeen = response.getDetailsList().stream()
            .map(s -> s.getDetails().getCell())
            .collect(Collectors.toSet());
    assertThat(cellsSeen).contains("one", "two");
}
 

@Test
public void allCellsOK() {
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(ok("one")));
    cellTwo.getServiceRegistry().addService(new CellWithHealthStatus(ok("two")));

    AssertableSubscriber<HealthCheckResponse> subscriber = service.check(HealthCheckRequest.newBuilder().build()).test();
    subscriber.awaitTerminalEvent(10, TimeUnit.SECONDS);
    subscriber.assertNoErrors();
    subscriber.assertValueCount(1);
    HealthCheckResponse response = subscriber.getOnNextEvents().get(0);
    assertThat(response.getStatus()).isEqualTo(SERVING);
    assertThat(response.getDetailsCount()).isEqualTo(2);
    assertThat(response.getDetails(0).hasDetails()).isTrue();
    assertThat(response.getDetails(1).hasDetails()).isTrue();
    Set<String> cellsSeen = response.getDetailsList().stream()
            .map(s -> s.getDetails().getCell())
            .collect(Collectors.toSet());
    assertThat(cellsSeen).contains("one", "two");
}
 

@Test
public void getPoliciesFromTwoCellsWithOneFailing() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithFailingAutoscalingService badCell = new CellWithFailingAutoscalingService();

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

    final AssertableSubscriber<GetPolicyResult> testSubscriber = service.getAllScalingPolicies(JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoValues();
    testSubscriber.assertError(StatusRuntimeException.class);
    List<Throwable> onErrorEvents = testSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents.size()).isEqualTo(1);
}
 

@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 observeJobDoesNotSetDeadlines() throws InterruptedException {
    String cellOneJobId = UUID.randomUUID().toString();
    String cellTwoJobId = UUID.randomUUID().toString();
    cellOne.getServiceRegistry().addService(new CellWithJobStream(cellOneJobId, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithJobStream(cellTwoJobId, cellTwoUpdates.serialize()));

    List<AssertableSubscriber<JobChangeNotification>> subscribers = new LinkedList<>();
    subscribers.add(service.observeJob(cellOneJobId, UNDEFINED_CALL_METADATA).test());
    subscribers.add(service.observeJob(cellTwoJobId, UNDEFINED_CALL_METADATA).test());

    // TODO: make it easier to extract the Deadline for each cell call
    Thread.sleep(2 * GRPC_REQUEST_TIMEOUT_MS);

    for (AssertableSubscriber<JobChangeNotification> subscriber : subscribers) {
        subscriber.assertNoErrors();
        subscriber.assertNotCompleted();
        assertThat(subscriber.isUnsubscribed()).isFalse();
    }
}
 

@Test
public void observeJobsErrorsWhenAnyClientsError() {
    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize()));

    final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test();

    final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-100").setStatus(ACCEPTED_STATE).build());
    final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-200").setStatus(ACCEPTED_STATE).build());
    cellOneUpdates.onNext(cellOneUpdate);
    cellTwoUpdates.onNext(cellTwoUpdate);

    testSubscriber.awaitValueCount(2, 1, TimeUnit.SECONDS);
    assertThat(testSubscriber.getOnErrorEvents()).isEmpty();
    assertThat(testSubscriber.isUnsubscribed()).isFalse();
    assertThat(testSubscriber.getCompletions()).isEqualTo(0);

    // a client emits an error
    cellTwoUpdates.onError(new RuntimeException("unexpected error"));

    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    assertThat(testSubscriber.getOnErrorEvents()).hasSize(1);
    assertThat(testSubscriber.isUnsubscribed()).isTrue();
    assertThat(testSubscriber.getCompletions()).isEqualTo(0);
}
 

@Test
public void nonRacy() throws Exception {
    Worker worker = Schedulers.computation().createWorker();
    try {
        for (int i = 0; i < 1000; i++) {
            AtomicInteger wip = new AtomicInteger(2);
    
            Subject<Integer, Integer> ps = PublishSubject.<Integer>create().toSerialized();
            
            AssertableSubscriber<Integer> as = ps.test(1);
            
            CountDownLatch cdl = new CountDownLatch(1);
            
            worker.schedule(() -> {
                if (wip.decrementAndGet() != 0) {
                    while (wip.get() != 0) ;
                }
                ps.onNext(1);
                
                cdl.countDown();
            });
            if (wip.decrementAndGet() != 0) {
                while (wip.get() != 0) ;
            }
            ps.onNext(1);
            
            cdl.await();
            
            as.assertFailure(MissingBackpressureException.class, 1);
        }
    } finally {
        worker.unsubscribe();
    }
}
 

@Test
public void observeJob() {
    String cellOneJobId = UUID.randomUUID().toString();
    String cellTwoJobId = UUID.randomUUID().toString();
    cellOne.getServiceRegistry().addService(new CellWithJobStream(cellOneJobId, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithJobStream(cellTwoJobId, cellTwoUpdates.serialize()));

    AssertableSubscriber<JobChangeNotification> subscriber1 = service.observeJob(cellOneJobId, UNDEFINED_CALL_METADATA).test();
    AssertableSubscriber<JobChangeNotification> subscriber2 = service.observeJob(cellTwoJobId, UNDEFINED_CALL_METADATA).test();

    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(ACCEPTED_STATE).build()));
    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(KILL_INITIATED_STATE).build()));
    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(FINISHED_STATE).build()));
    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(ACCEPTED_STATE).build()));

    subscriber1.awaitValueCount(3, 5, TimeUnit.SECONDS);
    subscriber1.assertNoErrors();
    subscriber1.assertNotCompleted();
    assertThat(subscriber1.isUnsubscribed()).isFalse();

    subscriber2.assertNoErrors();
    subscriber2.assertNoValues();
    subscriber2.assertNotCompleted();

    cellTwoUpdates.onNext(toNotification(Task.newBuilder()
            .setId(cellTwoJobId + "-task1").setJobId(cellTwoJobId)
            .build())
    );
    subscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);
    subscriber2.assertNoErrors();
    subscriber2.assertNotCompleted();

    cellOneUpdates.onCompleted();

    subscriber1.awaitTerminalEvent(1, TimeUnit.SECONDS);
    assertThat(subscriber1.getOnErrorEvents()).isEmpty();
    assertThat(subscriber1.isUnsubscribed()).isTrue();
    assertThat(subscriber1.getCompletions()).isEqualTo(1);
}
 

@Test
public void getJobLoadBalancersNoAssociations() {
    cellOne.getServiceRegistry().addService(new CellWithLoadBalancers(Collections.emptyList()));
    cellTwo.getServiceRegistry().addService(new CellWithLoadBalancers(Collections.emptyList()));

    final AssertableSubscriber<GetJobLoadBalancersResult> resultSubscriber = service.getLoadBalancers(
            JobId.newBuilder().setId(JOB_1).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors().assertCompleted();
    resultSubscriber.assertValueCount(1);
    assertThat(resultSubscriber.getOnNextEvents().get(0).getLoadBalancersCount()).isZero();
}
 

@Test
public void effectPerformerInvokesFunctionWithReceivedEffectAndEmitsReturnedEvents() {
  PublishSubject<String> upstream = PublishSubject.create();
  TestScheduler scheduler = new TestScheduler();
  Function<String, Integer> function = s -> s.length();
  AssertableSubscriber<Integer> observer =
      upstream.compose(Transformers.fromFunction(function, scheduler)).test();

  upstream.onNext("Hello");
  scheduler.triggerActions();
  observer.assertValue(5);
}
 

@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 observeJobsDoesNotSetDeadlines() throws InterruptedException {
    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize()));

    AssertableSubscriber<JobChangeNotification> subscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test();

    // TODO: make it easier to extract the Deadline for each cell call
    Thread.sleep(2 * GRPC_REQUEST_TIMEOUT_MS);

    subscriber.assertNoErrors();
    subscriber.assertNotCompleted();
    assertThat(subscriber.isUnsubscribed()).isFalse();
}