CassandraLoadBalancerStore implements LoadBalancerStore { @Override public Completable removeLoadBalancer(JobLoadBalancer jobLoadBalancer) { logger.debug("Removing load balancer {}", jobLoadBalancer); BoundStatement stmt = deleteAssociation.bind(jobLoadBalancer.getJobId(), jobLoadBalancer.getLoadBalancerId()); return storeHelper.execute(stmt) .map(rs -> { loadBalancerStateMap.remove(jobLoadBalancer); removeJobLoadBalancerAssociation(jobLoadBalancer); return jobLoadBalancer; }) .toCompletable(); } @Inject  CassandraLoadBalancerStore(CassandraStoreConfiguration configuration,
                                      @Named(LOAD_BALANCER_SANITIZER) EntitySanitizer entitySanitizer,
                                      @Named(LOAD_BALANCER_CASSANDRA_SESSION) Session session); @Activator void init(); @Override Observable<JobLoadBalancer> getAssociatedLoadBalancersForJob(String jobId); @Override Set<JobLoadBalancer> getAssociatedLoadBalancersSetForJob(String jobId); @Override List<JobLoadBalancerState> getAssociations(); @Override List<JobLoadBalancer> getAssociationsPage(int offset, int limit); @Override Completable addOrUpdateLoadBalancer(JobLoadBalancer jobLoadBalancer, JobLoadBalancer.State state); @Override Completable removeLoadBalancer(JobLoadBalancer jobLoadBalancer); @Override int getNumLoadBalancersForJob(String jobId); @Override Mono<Void> addOrUpdateTargets(Collection<LoadBalancerTargetState> targets); @Override Mono<Void> removeDeregisteredTargets(Collection<LoadBalancerTarget> toRemove); @Override Flux<LoadBalancerTargetState> getLoadBalancerTargets(String loadBalancerId); static final String LOAD_BALANCER_CASSANDRA_SESSION; }
@Test public void testRemoveLoadBalancer() throws Exception { Map<JobLoadBalancer, JobLoadBalancer.State> testData = generateTestData(10, 20, 1).getAssociations(); CassandraLoadBalancerStore store = getInitdStore(); testData.forEach((jobLoadBalancer, state) -> assertThat(store.addOrUpdateLoadBalancer(jobLoadBalancer, state).await(TIMEOUT_MS, TimeUnit.MILLISECONDS)).isTrue() ); testData.forEach((jobLoadBalancer, state) -> assertThat(store.removeLoadBalancer(jobLoadBalancer).await(TIMEOUT_MS, TimeUnit.MILLISECONDS)).isTrue() ); Map<String, List<JobLoadBalancerState>> byJobId = store.getAssociations().stream() .collect(Collectors.groupingBy(JobLoadBalancerState::getJobId)); testData.forEach((jobLoadBalancer, state) -> { assertThat(byJobId.get(jobLoadBalancer.getJobId())).isNullOrEmpty(); }); }