public <T> CompletableFuture<PushNotificationExperimentSample<T>> recordFinalState(final UUID accountIdentifier,
      final byte deviceId,
      final String experimentName,
      final T finalState) {

    CompletableFuture<String> finalStateJsonFuture;

    // Process the final state JSON on the calling thread, but inside a CompletionStage so there's just one "channel"
    // for reporting JSON exceptions. The alternative is to `throw JsonProcessingException`, but then callers would have
    // to both catch the exception when calling this method and also watch the returned future for the same exception.
    try {
      finalStateJsonFuture =
          CompletableFuture.completedFuture(SystemMapper.jsonMapper().writeValueAsString(finalState));
    } catch (final JsonProcessingException e) {
      finalStateJsonFuture = CompletableFuture.failedFuture(e);
    }

    final AttributeValue aciAndDeviceIdAttributeValue = buildSortKey(accountIdentifier, deviceId);

    return finalStateJsonFuture.thenCompose(finalStateJson -> {
      return dynamoDbAsyncClient.updateItem(UpdateItemRequest.builder()
              .tableName(tableName)
              .key(Map.of(
                  KEY_EXPERIMENT_NAME, AttributeValue.fromS(experimentName),
                  ATTR_ACI_AND_DEVICE_ID, aciAndDeviceIdAttributeValue))
              // `UpdateItem` will, by default, create a new item if one does not already exist for the given primary key. We
              // want update-only-if-exists behavior, though, and so check that there's already an existing item for this ACI
              // and device ID.
              .conditionExpression("#aciAndDeviceId = :aciAndDeviceId")
              .updateExpression("SET #finalState = if_not_exists(#finalState, :finalState)")
              .expressionAttributeNames(Map.of(
                  "#aciAndDeviceId", ATTR_ACI_AND_DEVICE_ID,
                  "#finalState", ATTR_FINAL_STATE))
              .expressionAttributeValues(Map.of(
                  ":aciAndDeviceId", aciAndDeviceIdAttributeValue,
                  ":finalState", AttributeValue.fromS(finalStateJson)))
              .returnValues(ReturnValue.ALL_NEW)
              .build())
          .thenApply(updateItemResponse -> {
            try {
              final boolean inExperimentGroup = updateItemResponse.attributes().get(ATTR_IN_EXPERIMENT_GROUP).bool();

              @SuppressWarnings("unchecked") final T parsedInitialState =
                  (T) parseState(updateItemResponse.attributes().get(ATTR_INITIAL_STATE).s(), finalState.getClass());

              @SuppressWarnings("unchecked") final T parsedFinalState =
                  (T) parseState(updateItemResponse.attributes().get(ATTR_FINAL_STATE).s(), finalState.getClass());

              return new PushNotificationExperimentSample<>(accountIdentifier, deviceId, inExperimentGroup, parsedInitialState, parsedFinalState);
            } catch (final JsonProcessingException e) {
              throw ExceptionUtils.wrap(e);
            }
          });
    });
  }
@Test
  void recordFinalState() throws JsonProcessingException {
    final String experimentName = "test-experiment";
    final UUID accountIdentifier = UUID.randomUUID();
    final byte deviceId = (byte) ThreadLocalRandom.current().nextInt(Device.MAXIMUM_DEVICE_ID);
    final boolean inExperimentGroup = ThreadLocalRandom.current().nextBoolean();
    final int initialBounciness = ThreadLocalRandom.current().nextInt();
    final int finalBounciness = initialBounciness + 17;

    {
      pushNotificationExperimentSamples.recordInitialState(accountIdentifier,
              deviceId,
              experimentName,
              inExperimentGroup,
              new TestDeviceState(initialBounciness))
          .join();

      final PushNotificationExperimentSample<TestDeviceState> returnedSample =
          pushNotificationExperimentSamples.recordFinalState(accountIdentifier,
                  deviceId,
                  experimentName,
                  new TestDeviceState(finalBounciness))
              .join();

      final PushNotificationExperimentSample<TestDeviceState> expectedSample =
          new PushNotificationExperimentSample<>(accountIdentifier, deviceId, inExperimentGroup,
          new TestDeviceState(initialBounciness),
          new TestDeviceState(finalBounciness));

      assertEquals(expectedSample, returnedSample,
          "Attempt to update existing sample without final state should succeed");

      assertEquals(expectedSample, getSample(accountIdentifier, deviceId, experimentName, TestDeviceState.class),
          "Attempt to update existing sample without final state should be persisted");
    }

    assertThrows(CompletionException.class, () -> pushNotificationExperimentSamples.recordFinalState(accountIdentifier,
            (byte) (deviceId + 1),
            experimentName,
            new TestDeviceState(finalBounciness))
        .join(),
        "Attempts to record a final state without an initial state should fail");
  }