public CompletableFuture<Integer> getCount(final UUID identifier, final byte deviceId) {
    final Timer.Sample sample = Timer.start();

    return dynamoDbAsyncClient.query(QueryRequest.builder()
            .tableName(tableName)
            .consistentRead(false)
            .keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
            .expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
            .expressionAttributeValues(Map.of(
                ":uuid", getPartitionKey(identifier),
                ":sortprefix", getSortKeyPrefix(deviceId)))
            .projectionExpression(ATTR_REMAINING_KEYS)
            .limit(1)
            .build())
        .thenApply(response -> {
          if (response.count() > 0) {
            final Map<String, AttributeValue> item = response.items().getFirst();

            if (item.containsKey(ATTR_REMAINING_KEYS)) {
              return Integer.parseInt(item.get(ATTR_REMAINING_KEYS).n());
            } else {
              // Some legacy keys sets may not have pre-counted keys; in that case, we'll tell the owners of those key
              // sets that they have none remaining, prompting an upload of a fresh set that we'll pre-count. This has
              // no effect on consumers of keys, which will still be able to take keys if any are actually present.
              noKeyCountAvailableCounter.increment();
              return 0;
            }
          } else {
            return 0;
          }
        })
        .whenComplete((keyCount, throwable) -> {
          sample.stop(getKeyCountTimer);

          if (throwable == null && keyCount != null) {
            availableKeyCountDistributionSummary.record(keyCount);
          }
        });
  }
@Test
  void getCount() {
    final SingleUsePreKeyStore<K> preKeyStore = getPreKeyStore();

    final UUID accountIdentifier = UUID.randomUUID();
    final byte deviceId = 1;

    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());

    final List<K> preKeys = generateRandomPreKeys();

    preKeyStore.store(accountIdentifier, deviceId, preKeys).join();

    assertEquals(KEY_COUNT, preKeyStore.getCount(accountIdentifier, deviceId).join());

    for (int i = 0; i < KEY_COUNT; i++) {
      preKeyStore.take(accountIdentifier, deviceId).join();
      assertEquals(KEY_COUNT - (i + 1), preKeyStore.getCount(accountIdentifier, deviceId).join());
    }

    preKeyStore.store(accountIdentifier, deviceId, List.of(generatePreKey(KEY_COUNT + 1))).join();
    clearKeyCountAttributes();

    assertEquals(0, preKeyStore.getCount(accountIdentifier, deviceId).join());
  }