@Override
  protected void run(final Environment environment,
      final Namespace namespace,
      final WhisperServerConfiguration configuration,
      final CommandDependencies commandDependencies) throws Exception {

    final PushNotificationExperiment<T> experiment =
        experimentFactory.buildExperiment(commandDependencies, configuration);

    final int maxConcurrency = namespace.getInt(MAX_CONCURRENCY_ARGUMENT);

    log.info("Finishing \"{}\" with max concurrency: {}", experiment.getExperimentName(), maxConcurrency);

    final AccountsManager accountsManager = commandDependencies.accountsManager();
    final PushNotificationExperimentSamples pushNotificationExperimentSamples = commandDependencies.pushNotificationExperimentSamples();

    final Flux<PushNotificationExperimentSample<T>> finishedSamples =
        pushNotificationExperimentSamples.getSamples(experiment.getExperimentName(), experiment.getStateClass())
            .doOnNext(sample -> Metrics.counter(SAMPLES_READ_COUNTER_NAME, "final", String.valueOf(sample.finalState() != null)).increment())
            .flatMap(sample -> {
              if (sample.finalState() == null) {
                // We still need to record a final state for this sample
                return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(sample.accountIdentifier()))
                    .retryWhen(Retry.backoff(3, Duration.ofSeconds(1)))
                    .doOnNext(ignored -> ACCOUNT_READ_COUNTER.increment())
                    .flatMap(maybeAccount -> {
                      final T finalState = experiment.getState(maybeAccount.orElse(null),
                          maybeAccount.flatMap(account -> account.getDevice(sample.deviceId())).orElse(null));

                      return Mono.fromFuture(
                              () -> pushNotificationExperimentSamples.recordFinalState(sample.accountIdentifier(),
                                  sample.deviceId(),
                                  experiment.getExperimentName(),
                                  finalState))
                          .onErrorResume(ConditionalCheckFailedException.class, throwable -> Mono.empty())
                          .onErrorResume(JsonProcessingException.class, throwable -> {
                            log.error("Failed to parse sample state JSON", throwable);
                            return Mono.empty();
                          })
                          .retryWhen(Retry.backoff(3, Duration.ofSeconds(1)))
                          .onErrorResume(throwable -> {
                            log.warn("Failed to record final state for {}:{} in experiment {}",
                                sample.accountIdentifier(), sample.deviceId(), experiment.getExperimentName(), throwable);

                            return Mono.empty();
                          })
                          .doOnSuccess(ignored -> FINAL_SAMPLE_STORED_COUNTER.increment());
                    });
              } else {
                return Mono.just(sample);
              }
            }, maxConcurrency);

    experiment.analyzeResults(finishedSamples);
  }
@Test
  void runStoreSampleRetry() {
    final UUID accountIdentifier = UUID.randomUUID();
    final byte deviceId = Device.PRIMARY_ID;

    final Device device = mock(Device.class);
    when(device.getId()).thenReturn(deviceId);

    final Account account = mock(Account.class);
    when(account.getDevice(deviceId)).thenReturn(Optional.of(device));

    when(commandDependencies.accountsManager().getByAccountIdentifierAsync(accountIdentifier))
        .thenReturn(CompletableFuture.completedFuture(Optional.of(account)));

    when(commandDependencies.pushNotificationExperimentSamples().getSamples(eq(EXPERIMENT_NAME), eq(String.class)))
        .thenReturn(Flux.just(new PushNotificationExperimentSample<>(accountIdentifier, deviceId, true, "test", null)));

    when(commandDependencies.pushNotificationExperimentSamples().recordFinalState(any(), anyByte(), any(), any()))
        .thenReturn(CompletableFuture.failedFuture(new RuntimeException()))
        .thenReturn(CompletableFuture.failedFuture(new RuntimeException()))
        .thenReturn(CompletableFuture.completedFuture(new PushNotificationExperimentSample<>(accountIdentifier, deviceId, true, "test", "test")));

    assertDoesNotThrow(() -> finishPushNotificationExperimentCommand.run(null, NAMESPACE, null, commandDependencies));
    verify(experiment).getState(account, device);
    verify(commandDependencies.pushNotificationExperimentSamples(), times(3))
        .recordFinalState(eq(accountIdentifier), eq(deviceId), eq(EXPERIMENT_NAME), any());
  }