software.amazon.awssdk.services.dynamodb.model.GetItemRequest Java Examples

@Test
public void getApplication() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  GetItemRequest expectedGetItemRequest = GetItemRequest.builder()
        .tableName(TABLE_NAME)
        .consistentRead(Boolean.TRUE)
        .key(recordMap)
        .build();

  Application application = service.getApplication(applicationId);
  ArgumentCaptor<GetItemRequest> getItemRequestArgumentCaptor = ArgumentCaptor.forClass(GetItemRequest.class);
  verify(dynamodb).getItem(getItemRequestArgumentCaptor.capture());

  assertThat(application.getApplicationId()).isEqualTo(applicationId);
  assertThat(getItemRequestArgumentCaptor.getValue()).isEqualTo(expectedGetItemRequest);
}
 

@Test
public void generateRequest_partitionAndSortKey() {
    FakeItemWithSort keyItem = createUniqueFakeItemWithSort();
    GetItemOperation<FakeItemWithSort> getItemOperation =
        GetItemOperation.create(GetItemEnhancedRequest.builder()
                                                      .key(k -> k.partitionValue(keyItem.getId())
                                                                 .sortValue(keyItem.getSort()))
                                                      .build());

    GetItemRequest request = getItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                              PRIMARY_CONTEXT,
                                                              null);

    Map<String, AttributeValue> expectedKeyMap = new HashMap<>();
    expectedKeyMap.put("id", AttributeValue.builder().s(keyItem.getId()).build());
    expectedKeyMap.put("sort", AttributeValue.builder().s(keyItem.getSort()).build());
    GetItemRequest expectedRequest = GetItemRequest.builder()
                                                   .tableName(TABLE_NAME)
                                                   .key(expectedKeyMap)
                                                   .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void generateRequest_partitionKeyOnly() {
    FakeItem keyItem = createUniqueFakeItem();
    GetItemOperation<FakeItem> getItemOperation =
        GetItemOperation.create(GetItemEnhancedRequest.builder().key(k -> k.partitionValue(keyItem.getId())).build());

    GetItemRequest request = getItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                              PRIMARY_CONTEXT,
                                                              null);

    Map<String, AttributeValue> expectedKeyMap = new HashMap<>();
    expectedKeyMap.put("id", AttributeValue.builder().s(keyItem.getId()).build());
    GetItemRequest expectedRequest = GetItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(expectedKeyMap)
        .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void generateRequest_consistentRead() {
    FakeItem keyItem = createUniqueFakeItem();
    GetItemOperation<FakeItem> getItemOperation =
        GetItemOperation.create(GetItemEnhancedRequest.builder()
                                                      .key(k -> k.partitionValue(keyItem.getId()))
                                                      .consistentRead(true).build());

    GetItemRequest request = getItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                              PRIMARY_CONTEXT,
                                                              null);

    Map<String, AttributeValue> expectedKeyMap = new HashMap<>();
    expectedKeyMap.put("id", AttributeValue.builder().s(keyItem.getId()).build());
    GetItemRequest expectedRequest = GetItemRequest.builder()
                                                   .tableName(TABLE_NAME)
                                                   .key(expectedKeyMap)
                                                   .consistentRead(true)
                                                   .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void testGetJob_Failure() throws Throwable {
  String token = getToken(4);
  JobTokenPayload tokenPayload = tokenManager.decodeToken(token);
  Map<String, AttributeValue> primaryKey =
      DynamoDBUtils.getPrimaryKeyFromToken(token, tokenPayload, maxShardId);

  GetItemRequest getItemRequest =
      GetItemRequest.builder()
          .key(primaryKey)
          .tableName(tableName)
          .attributesToGet(Collections.singletonList(DynamoDBUtils.Attribute.JOB_SPEC.columnName))
          .build();

  CompletableFuture<GetItemResponse> ret = new CompletableFuture<>();
  Exception exception = new Exception();
  ret.completeExceptionally(exception);
  when(ddbClient.getItem(getItemRequest)).thenReturn(ret);

  CompletableFuture<Schedule> response = scheduleManager.getJob(token);
  assertSame(getException(response), exception);

  verify(ddbClient, times(1)).getItem(getItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

@Override
public CompletionStage<Map<Integer, Long>> findAll(String topic, GroupId groupId) {
    var request = GetItemRequest.builder()
            .tableName(tableName)
            .consistentRead(true)
            .key(toKey(topic, groupId))
            .build();

    return Mono.fromCompletionStage(() -> dynamoDB.getItem(request))
            .<Map<Integer, Long>>handle((result, sink) -> {
                try {
                    var positions = toPositions(result.item());

                    if (positions == null) {
                        sink.complete();
                    } else {
                        sink.next(positions);
                    }
                } catch (Exception e) {
                    sink.error(e);
                }
            })
            .log(this.getClass().getName(), Level.WARNING, SignalType.ON_ERROR)
            .retryWhen(it -> it.delayElements(Duration.ofSeconds(1)))
            .toFuture();
}
 

@Test
public void getEmptyString() {
    Map<String, AttributeValue> itemMap = new HashMap<>();
    itemMap.put("id", AttributeValue.builder().s("id123").build());
    itemMap.put("s", EMPTY_STRING);

    GetItemResponse response = GetItemResponse.builder()
                                              .item(itemMap)
                                              .build();

    when(mockDynamoDbClient.getItem(any(GetItemRequest.class))).thenReturn(response);

    TestBean result = dynamoDbTable.getItem(r -> r.key(k -> k.partitionValue("id123")));

    assertThat(result.getId()).isEqualTo("id123");
    assertThat(result.getS()).isEmpty();
}
 

@Test
public void getEmptyBytes() {
    Map<String, AttributeValue> itemMap = new HashMap<>();
    itemMap.put("id", AttributeValue.builder().s("id123").build());
    itemMap.put("b", EMPTY_BINARY);

    GetItemResponse response = GetItemResponse.builder()
                                              .item(itemMap)
                                              .build();

    when(mockDynamoDbClient.getItem(any(GetItemRequest.class))).thenReturn(response);

    TestBean result = dynamoDbTable.getItem(r -> r.key(k -> k.partitionValue("id123")));

    assertThat(result.getId()).isEqualTo("id123");
    assertThat(result.getB()).isEqualTo(EMPTY_BYTES);
}
 

@Test
public void testGetLeaseTimesOut() throws Exception {
    TimeoutException te = setRuleForDependencyTimeout();

    when(dynamoDbClient.getItem(any(GetItemRequest.class))).thenReturn(mockGetItemFuture);
    when(mockGetItemFuture.get(anyLong(), any())).thenThrow(te);

    when(leaseSerializer.getDynamoHashKey(anyString())).thenReturn(Collections.emptyMap());

    verifyCancel(mockGetItemFuture, () -> leaseRefresher.getLease("test"));
}
 

private ApplicationRecord loadApplication(final String applicationId) {
  Map<String, AttributeValue> applicationMap = dynamodb.getItem(GetItemRequest.builder()
        .tableName(tableName)
        .consistentRead(Boolean.TRUE)
        .key(toKeyRecord(applicationId))
        .build()).item();
  if (applicationMap == null || applicationMap.isEmpty()) {
    throw new NotFoundApiException(new NotFoundException()
          .errorCode("ApplicationNotFound")
          .message(String.format("Application %s can not be found.", applicationId)));
  }
  return new ApplicationRecord(applicationMap);
}
 

private Map<String, AttributeValue> getLockItem(String lockName) {
    GetItemRequest request = GetItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(Collections.singletonMap(ID, AttributeValue.builder()
            .s(lockName)
            .build()))
        .build();
    GetItemResponse response = dynamodb.getItem(request);
    return response.item();
}
 

public static void getDynamoDBItem(DynamoDbClient ddb,String tableName,String key,String keyVal ) {

        HashMap<String,AttributeValue> keyToGet = new HashMap<String,AttributeValue>();

        keyToGet.put(key, AttributeValue.builder()
                .s(keyVal).build());

        // Create a GetItemRequest object
        GetItemRequest request = GetItemRequest.builder()
                .key(keyToGet)
                .tableName(tableName)
                .build();

        try {
            Map<String,AttributeValue> returnedItem = ddb.getItem(request).item();

            if (returnedItem != null) {
                Set<String> keys = returnedItem.keySet();
                System.out.println("Table Attributes: \n");

                for (String key1 : keys) {
                    System.out.format("%s: %s\n", key1, returnedItem.get(key1).toString());
                }
            } else {
                System.out.format("No item found with the key %s!\n", key);
            }
        } catch (DynamoDbException e) {
            System.err.println(e.getMessage());
            System.exit(1);
        }
        // snippet-end:[dynamodb.java2.get_item.main]
    }
 

@Test
public void ddbRequest_shouldModifyRequest() {
    GetItemRequest getItemRequest = GetItemRequest.builder().build();
    SdkRequest sdkRequest = interceptor.modifyRequest(() -> getItemRequest, new ExecutionAttributes());

    RequestOverrideConfiguration requestOverrideConfiguration = sdkRequest.overrideConfiguration().get();
    assertThat(requestOverrideConfiguration.apiNames()
                                           .stream()
                                           .filter(a -> a.name()
                                                         .equals("hll") &&
                                                        a.version().equals("ddb-enh")).findAny())
        .isPresent();
}
 

@Test
public void generateRequest_withExtension_doesNotModifyKey() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> keyMap = FakeItem.getTableSchema().itemToMap(baseFakeItem, singletonList("id"));
    GetItemOperation<FakeItem> getItemOperation =
        GetItemOperation.create(GetItemEnhancedRequest.builder().key(k -> k.partitionValue(baseFakeItem.getId())).build());

    GetItemRequest request = getItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                              PRIMARY_CONTEXT,
                                                              mockDynamoDbEnhancedClientExtension);

    assertThat(request.key(), is(keyMap));
    verify(mockDynamoDbEnhancedClientExtension, never()).beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class));
}
 

@Test
public void getServiceCall_makesTheRightCallAndReturnsResponse() {
    FakeItem keyItem = createUniqueFakeItem();
    GetItemOperation<FakeItem> getItemOperation =
        GetItemOperation.create(GetItemEnhancedRequest.builder().key(k -> k.partitionValue(keyItem.getId())).build());
    GetItemRequest getItemRequest = GetItemRequest.builder().tableName(TABLE_NAME).build();
    GetItemResponse expectedResponse = GetItemResponse.builder().build();
    when(mockDynamoDbClient.getItem(any(GetItemRequest.class))).thenReturn(expectedResponse);

    GetItemResponse response = getItemOperation.serviceCall(mockDynamoDbClient).apply(getItemRequest);

    assertThat(response, sameInstance(expectedResponse));
    verify(mockDynamoDbClient).getItem(getItemRequest);
}
 

@Override
public GetItemRequest generateRequest(TableSchema<T> tableSchema,
                                      OperationContext context,
                                      DynamoDbEnhancedClientExtension extension) {
    if (!TableMetadata.primaryIndexName().equals(context.indexName())) {
        throw new IllegalArgumentException("GetItem cannot be executed against a secondary index.");
    }

    return GetItemRequest.builder()
                         .tableName(context.tableName())
                         .key(this.request.key().keyMap(tableSchema, context.indexName()))
                         .consistentRead(this.request.consistentRead())
                         .build();
}
 

@Test
public void getApplication_null_notFound() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  GetItemResponse response = GetItemResponse.builder().build();
  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  assertThatThrownBy(() -> service.getApplication(applicationId))
        .isInstanceOf(NotFoundApiException.class);
}
 

@Test
public void getApplication_empty_notFound() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  GetItemResponse response = GetItemResponse.builder().item(new HashMap<>()).build();
  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  assertThatThrownBy(() -> service.getApplication(applicationId))
        .isInstanceOf(NotFoundApiException.class);
}
 

@Test
public void deleteApplication() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  Long version = 1L;
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  recordMap.put("version", AttributeValue.builder().n(version.toString()).build());
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  Map<String, AttributeValue> expectedAttributeValueMap = new HashMap<>();
  expectedAttributeValueMap.put(":v", AttributeValue.builder().n(version.toString()).build());
  DeleteItemRequest expectedDeleteItemRequest = DeleteItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(keyMap(userId, applicationId))
        .conditionExpression("version = :v")
        .expressionAttributeValues(expectedAttributeValueMap)
        .build();

  service.deleteApplication(applicationId);
  ArgumentCaptor<DeleteItemRequest> deleteItemRequestArgumentCaptor = ArgumentCaptor.forClass(DeleteItemRequest.class);
  verify(dynamodb).deleteItem(deleteItemRequestArgumentCaptor.capture());

  DeleteItemRequest deleteItemRequest = deleteItemRequestArgumentCaptor.getValue();
  assertThat(deleteItemRequest).isEqualTo(expectedDeleteItemRequest);
}
 

private void getItemAndAssertValues(DynamoDbClient client) {
    Map<String, AttributeValue> item =
        client.getItem(GetItemRequest.builder()
                                     .tableName(TABLE_NAME)
                                     .key(Collections.singletonMap(HASH_KEY_NAME, AttributeValue.builder()
                                                                                                .s(HASH_KEY_VALUE)
                                                                                                .build()))
                                     .build())
              .item();

    assertEquals(HASH_KEY_VALUE, item.get(HASH_KEY_NAME).s());
    assertEquals(ATTRIBUTE_FOO_VALUE, item.get(ATTRIBUTE_FOO).s());
}
 

@Test
public void updateApplication_author() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  String author = UUID.randomUUID().toString();
  Long version = 1L;
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  recordMap.put("version", AttributeValue.builder().n(version.toString()).build());
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  Map<String, AttributeValue> updateMap = new HashMap<>();
  updateMap.put(":nv", AttributeValue.builder().n("2").build());
  updateMap.put(":v", AttributeValue.builder().n("1").build());
  updateMap.put(":a", AttributeValue.builder().s(author).build());

  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  UpdateApplicationInput input = new UpdateApplicationInput().author(author);
  UpdateItemRequest expectedUpdateItemRequest = UpdateItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(keyMap(userId, applicationId))
        .updateExpression("SET author = :a,version = :nv")
        .expressionAttributeValues(updateMap)
        .conditionExpression("version = :v")
        .build();

  Application application = service.updateApplication(input, applicationId);
  ArgumentCaptor<UpdateItemRequest> updateItemRequestArgumentCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class);
  verify(dynamodb).updateItem(updateItemRequestArgumentCaptor.capture());

  UpdateItemRequest updateItemRequest = updateItemRequestArgumentCaptor.getValue();
  assertThat(updateItemRequest).isEqualTo(expectedUpdateItemRequest);
  assertThat(application.getApplicationId()).isEqualTo(applicationId);
  assertThat(application.getAuthor()).isEqualTo(author);
}
 

@Test
public void updateApplication_description() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  String description = UUID.randomUUID().toString();
  Long version = 1L;
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  recordMap.put("version", AttributeValue.builder().n(version.toString()).build());
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  Map<String, AttributeValue> updateMap = new HashMap<>();
  updateMap.put(":nv", AttributeValue.builder().n("2").build());
  updateMap.put(":v", AttributeValue.builder().n("1").build());
  updateMap.put(":d", AttributeValue.builder().s(description).build());

  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  UpdateApplicationInput input = new UpdateApplicationInput().description(description);
  UpdateItemRequest expectedUpdateItemRequest = UpdateItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(keyMap(userId, applicationId))
        .updateExpression("SET description = :d,version = :nv")
        .expressionAttributeValues(updateMap)
        .conditionExpression("version = :v")
        .build();

  Application application = service.updateApplication(input, applicationId);
  ArgumentCaptor<UpdateItemRequest> updateItemRequestArgumentCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class);
  verify(dynamodb).updateItem(updateItemRequestArgumentCaptor.capture());

  UpdateItemRequest updateItemRequest = updateItemRequestArgumentCaptor.getValue();
  assertThat(updateItemRequest).isEqualTo(expectedUpdateItemRequest);
  assertThat(application.getApplicationId()).isEqualTo(applicationId);
  assertThat(application.getDescription()).isEqualTo(description);
}
 

@Test
public void updateApplication_homePageUrl() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  String homePageUrl = UUID.randomUUID().toString();
  Long version = 1L;
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  recordMap.put("version", AttributeValue.builder().n(version.toString()).build());
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  Map<String, AttributeValue> updateMap = new HashMap<>();
  updateMap.put(":nv", AttributeValue.builder().n("2").build());
  updateMap.put(":v", AttributeValue.builder().n("1").build());
  updateMap.put(":h", AttributeValue.builder().s(homePageUrl).build());

  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  UpdateApplicationInput input = new UpdateApplicationInput().homePageUrl(homePageUrl);
  UpdateItemRequest expectedUpdateItemRequest = UpdateItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(keyMap(userId, applicationId))
        .updateExpression("SET homePageUrl = :h,version = :nv")
        .expressionAttributeValues(updateMap)
        .conditionExpression("version = :v")
        .build();

  Application application = service.updateApplication(input, applicationId);
  ArgumentCaptor<UpdateItemRequest> updateItemRequestArgumentCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class);
  verify(dynamodb).updateItem(updateItemRequestArgumentCaptor.capture());

  UpdateItemRequest updateItemRequest = updateItemRequestArgumentCaptor.getValue();
  assertThat(updateItemRequest).isEqualTo(expectedUpdateItemRequest);
  assertThat(application.getApplicationId()).isEqualTo(applicationId);
  assertThat(application.getHomePageUrl()).isEqualTo(homePageUrl);
}
 

@Override
public CompletableFuture<Schedule> getJob(String token) {
  try {
    JobTokenPayload tokenPayload = tokenManager.decodeToken(token);
    Map<String, AttributeValue> primaryKey =
        DynamoDBUtils.getPrimaryKeyFromToken(token, tokenPayload, maxShardId);

    GetItemRequest getItemRequest =
        GetItemRequest.builder()
            .key(primaryKey)
            .tableName(ddbConfig.getSchedulesTableName())
            .attributesToGet(Collections.singletonList(Attribute.JOB_SPEC.columnName))
            .build();

    return ddbClient
        .getItem(getItemRequest)
        .thenApply(
            item -> {
              try {
                return makeSchedule(item.item().get(Attribute.JOB_SPEC.columnName).s());
              } catch (InvalidTokenException e) {
                throw new RuntimeException(e);
              }
            });
  } catch (InvalidTokenException ex) {
    CompletableFuture<Schedule> future = new CompletableFuture<>();
    future.completeExceptionally(ex);
    return future;
  }
}
 

@Test
public void testGetJob() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);

  JobTokenPayload tokenPayload = tokenManager.decodeToken(validToken);
  Map<String, AttributeValue> primaryKey =
      DynamoDBUtils.getPrimaryKeyFromToken(validToken, tokenPayload, maxShardId);

  GetItemRequest getItemRequest =
      GetItemRequest.builder()
          .key(primaryKey)
          .tableName(tableName)
          .attributesToGet(Collections.singletonList(DynamoDBUtils.Attribute.JOB_SPEC.columnName))
          .build();

  when(ddbClient.getItem(getItemRequest))
      .thenReturn(
          CompletableFuture.completedFuture(
              GetItemResponse.builder()
                  .item(
                      ImmutableMap.of(
                          DynamoDBUtils.Attribute.JOB_SPEC.columnName,
                          AttributeValue.builder().s(schedule.getJobSpec()).build()))
                  .build()));

  CompletableFuture<Schedule> response = scheduleManager.getJob(validToken);

  Assert.assertEquals(response.get(1000, TimeUnit.MILLISECONDS), schedule);

  verify(ddbClient, times(1)).getItem(getItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

protected GetItemRequest getRequest(String name) {
    Map<String, AttributeValue> key = new HashMap<>();
    key.put(FRUIT_NAME_COL, AttributeValue.builder().s(name).build());

    return GetItemRequest.builder()
            .tableName(getTableName())
            .key(key)
            .attributesToGet(FRUIT_NAME_COL, FRUIT_DESC_COL)
            .build();
}
 

public static GetItemRequest createGetRequest(String table, String keyValue, String rangeValue) {
    Map<String, AttributeValue> key = new HashMap<>();
    key.put(KEY_NAME, AttributeValue.builder().s(keyValue).build());
    key.put(RANGE_NAME, AttributeValue.builder().s(rangeValue).build());

    return GetItemRequest.builder()
            .tableName(table)
            .key(key)
            .attributesToGet(PAYLOAD_NAME)
            .build();
}
 

@Test
void normal() {
  server.enqueue(HttpResponse.of(HttpStatus.OK));

  GetItemResponse response =
      dynamoDbAsync
          .getItem(GetItemRequest.builder().tableName("test").key(ImmutableMap.of()).build())
          .join();
  assertThat(response.sdkHttpResponse().isSuccessful()).isTrue();

  AggregatedHttpRequest request = server.takeRequest().request();
  assertThat(request.headers().get(HttpHeaderNames.USER_AGENT)).contains("http/ArmeriaAsync");
}
 

@Test
public void twoVersionsWithTwoMaterialsWithRefresh() {
    final Map<String, AttributeValue> attr1 = Collections.singletonMap(MATERIAL_PARAM, AttributeValueBuilder.ofS("material1"));
    final EncryptionContext ctx1 = ctx(attr1);
    final Map<String, AttributeValue> attr2 = Collections.singletonMap(MATERIAL_PARAM, AttributeValueBuilder.ofS("material2"));
    final EncryptionContext ctx2 = ctx(attr2);

    final MostRecentProvider prov = new ExtendedProvider(store, 500);
    verify(client, never()).putItem(any(PutItemRequest.class));
    final EncryptionMaterials eMat1_1 = prov.getEncryptionMaterials(ctx1);
    // It's a new provider, so we see a single putItem
    verify(client).putItem(any(PutItemRequest.class));
    reset(client);
    final EncryptionMaterials eMat1_2 = prov.getEncryptionMaterials(ctx2);
    // It's a new provider, so we see a single putItem
    verify(client).putItem(any(PutItemRequest.class));
    reset(client);
    // Create the new material
    store.newProvider("material1");
    store.newProvider("material2");
    reset(client);
    // Ensure the cache is working
    final EncryptionMaterials eMat2_1 = prov.getEncryptionMaterials(ctx1);
    final EncryptionMaterials eMat2_2 = prov.getEncryptionMaterials(ctx2);
    verifyNoMoreInteractions(client);
    assertEquals(0, store.getVersionFromMaterialDescription(eMat1_1.getMaterialDescription()));
    assertEquals(0, store.getVersionFromMaterialDescription(eMat2_1.getMaterialDescription()));
    assertEquals(0, store.getVersionFromMaterialDescription(eMat1_2.getMaterialDescription()));
    assertEquals(0, store.getVersionFromMaterialDescription(eMat2_2.getMaterialDescription()));
    prov.refresh();
    final EncryptionMaterials eMat3_1 = prov.getEncryptionMaterials(ctx1);
    final EncryptionMaterials eMat3_2 = prov.getEncryptionMaterials(ctx2);
    verify(client, times(2)).query(any(QueryRequest.class)); // To find current version
    verify(client, times(2)).getItem(any(GetItemRequest.class));
    assertEquals(1, store.getVersionFromMaterialDescription(eMat3_1.getMaterialDescription()));
    assertEquals(1, store.getVersionFromMaterialDescription(eMat3_2.getMaterialDescription()));

    assertEquals(eMat1_1.getSigningKey(), eMat2_1.getSigningKey());
    assertFalse(eMat1_1.getSigningKey().equals(eMat3_1.getSigningKey()));
    assertEquals(eMat1_2.getSigningKey(), eMat2_2.getSigningKey());
    assertFalse(eMat1_2.getSigningKey().equals(eMat3_2.getSigningKey()));

    // Ensure we can decrypt all of them without hitting ddb more than the minimum
    final MostRecentProvider prov2 = new ExtendedProvider(store, 500);
    final DecryptionMaterials dMat1_1 = prov2.getDecryptionMaterials(ctx(eMat1_1, attr1));
    final DecryptionMaterials dMat1_2 = prov2.getDecryptionMaterials(ctx(eMat1_2, attr2));
    reset(client);
    assertEquals(eMat1_1.getEncryptionKey(), dMat1_1.getDecryptionKey());
    assertEquals(eMat1_2.getEncryptionKey(), dMat1_2.getDecryptionKey());
    assertEquals(eMat1_1.getSigningKey(), dMat1_1.getVerificationKey());
    assertEquals(eMat1_2.getSigningKey(), dMat1_2.getVerificationKey());
    final DecryptionMaterials dMat2_1 = prov2.getDecryptionMaterials(ctx(eMat2_1, attr1));
    final DecryptionMaterials dMat2_2 = prov2.getDecryptionMaterials(ctx(eMat2_2, attr2));
    assertEquals(eMat2_1.getEncryptionKey(), dMat2_1.getDecryptionKey());
    assertEquals(eMat2_2.getEncryptionKey(), dMat2_2.getDecryptionKey());
    assertEquals(eMat2_1.getSigningKey(), dMat2_1.getVerificationKey());
    assertEquals(eMat2_2.getSigningKey(), dMat2_2.getVerificationKey());
    final DecryptionMaterials dMat3_1 = prov2.getDecryptionMaterials(ctx(eMat3_1, attr1));
    final DecryptionMaterials dMat3_2 = prov2.getDecryptionMaterials(ctx(eMat3_2, attr2));
    assertEquals(eMat3_1.getEncryptionKey(), dMat3_1.getDecryptionKey());
    assertEquals(eMat3_2.getEncryptionKey(), dMat3_2.getDecryptionKey());
    assertEquals(eMat3_1.getSigningKey(), dMat3_1.getVerificationKey());
    assertEquals(eMat3_2.getSigningKey(), dMat3_2.getVerificationKey());
    // Get item will be hit once for the one old key
    verify(client, times(2)).getItem(any(GetItemRequest.class));
    verifyNoMoreInteractions(client);
}
 

@Test
public void singleVersionWithTwoMaterialsWithRefresh() throws InterruptedException {
    final Map<String, AttributeValue> attr1 = Collections.singletonMap(MATERIAL_PARAM, AttributeValueBuilder.ofS("material1"));
    final EncryptionContext ctx1 = ctx(attr1);
    final Map<String, AttributeValue> attr2 = Collections.singletonMap(MATERIAL_PARAM, AttributeValueBuilder.ofS("material2"));
    final EncryptionContext ctx2 = ctx(attr2);

    final MostRecentProvider prov = new ExtendedProvider(store, 500);
    verify(client, never()).putItem(any(PutItemRequest.class));
    final EncryptionMaterials eMat1_1 = prov.getEncryptionMaterials(ctx1);
    // It's a new provider, so we see a single putItem
    verify(client).putItem(any(PutItemRequest.class));
    reset(client);
    final EncryptionMaterials eMat1_2 = prov.getEncryptionMaterials(ctx2);
    // It's a new provider, so we see a single putItem
    verify(client).putItem(any(PutItemRequest.class));
    reset(client);
    // Ensure the two materials are, in fact, different
    assertFalse(eMat1_1.getSigningKey().equals(eMat1_2.getSigningKey()));

    // Ensure the cache is working
    final EncryptionMaterials eMat2_1 = prov.getEncryptionMaterials(ctx1);
    verifyNoMoreInteractions(client);
    assertEquals(0, store.getVersionFromMaterialDescription(eMat1_1.getMaterialDescription()));
    assertEquals(0, store.getVersionFromMaterialDescription(eMat2_1.getMaterialDescription()));
    final EncryptionMaterials eMat2_2 = prov.getEncryptionMaterials(ctx2);
    verifyNoMoreInteractions(client);
    assertEquals(0, store.getVersionFromMaterialDescription(eMat1_2.getMaterialDescription()));
    assertEquals(0, store.getVersionFromMaterialDescription(eMat2_2.getMaterialDescription()));

    prov.refresh();
    final EncryptionMaterials eMat3_1 = prov.getEncryptionMaterials(ctx1);
    verify(client).query(any(QueryRequest.class)); // To find current version
    verify(client).getItem(any(GetItemRequest.class));
    final EncryptionMaterials eMat3_2 = prov.getEncryptionMaterials(ctx2);
    verify(client, times(2)).query(any(QueryRequest.class)); // To find current version
    verify(client, times(2)).getItem(any(GetItemRequest.class));
    assertEquals(0, store.getVersionFromMaterialDescription(eMat3_1.getMaterialDescription()));
    assertEquals(0, store.getVersionFromMaterialDescription(eMat3_2.getMaterialDescription()));
    prov.refresh();

    assertEquals(eMat1_1.getSigningKey(), eMat2_1.getSigningKey());
    assertEquals(eMat1_1.getSigningKey(), eMat3_1.getSigningKey());
    assertEquals(eMat1_2.getSigningKey(), eMat2_2.getSigningKey());
    assertEquals(eMat1_2.getSigningKey(), eMat3_2.getSigningKey());

    // Ensure that after cache refresh  we only get one more hit as opposed to multiple
    prov.getEncryptionMaterials(ctx1);
    prov.getEncryptionMaterials(ctx2);
    Thread.sleep(700);
    // Force refresh
    prov.getEncryptionMaterials(ctx1);
    prov.getEncryptionMaterials(ctx2);
    reset(client);
    // Check to ensure no more hits
    assertEquals(eMat1_1.getSigningKey(), prov.getEncryptionMaterials(ctx1).getSigningKey());
    assertEquals(eMat1_1.getSigningKey(), prov.getEncryptionMaterials(ctx1).getSigningKey());
    assertEquals(eMat1_1.getSigningKey(), prov.getEncryptionMaterials(ctx1).getSigningKey());
    assertEquals(eMat1_1.getSigningKey(), prov.getEncryptionMaterials(ctx1).getSigningKey());
    assertEquals(eMat1_1.getSigningKey(), prov.getEncryptionMaterials(ctx1).getSigningKey());

    assertEquals(eMat1_2.getSigningKey(), prov.getEncryptionMaterials(ctx2).getSigningKey());
    assertEquals(eMat1_2.getSigningKey(), prov.getEncryptionMaterials(ctx2).getSigningKey());
    assertEquals(eMat1_2.getSigningKey(), prov.getEncryptionMaterials(ctx2).getSigningKey());
    assertEquals(eMat1_2.getSigningKey(), prov.getEncryptionMaterials(ctx2).getSigningKey());
    assertEquals(eMat1_2.getSigningKey(), prov.getEncryptionMaterials(ctx2).getSigningKey());
    verifyNoMoreInteractions(client);

    // Ensure we can decrypt all of them without hitting ddb more than the minimum
    final MostRecentProvider prov2 = new ExtendedProvider(store, 500);
    final DecryptionMaterials dMat1_1 = prov2.getDecryptionMaterials(ctx(eMat1_1, attr1));
    final DecryptionMaterials dMat1_2 = prov2.getDecryptionMaterials(ctx(eMat1_2, attr2));
    reset(client);
    assertEquals(eMat1_1.getEncryptionKey(), dMat1_1.getDecryptionKey());
    assertEquals(eMat1_2.getEncryptionKey(), dMat1_2.getDecryptionKey());
    assertEquals(eMat1_1.getSigningKey(), dMat1_1.getVerificationKey());
    assertEquals(eMat1_2.getSigningKey(), dMat1_2.getVerificationKey());
    final DecryptionMaterials dMat2_1 = prov2.getDecryptionMaterials(ctx(eMat2_1, attr1));
    final DecryptionMaterials dMat2_2 = prov2.getDecryptionMaterials(ctx(eMat2_2, attr2));
    assertEquals(eMat2_1.getEncryptionKey(), dMat2_1.getDecryptionKey());
    assertEquals(eMat2_2.getEncryptionKey(), dMat2_2.getDecryptionKey());
    assertEquals(eMat2_1.getSigningKey(), dMat2_1.getVerificationKey());
    assertEquals(eMat2_2.getSigningKey(), dMat2_2.getVerificationKey());
    final DecryptionMaterials dMat3_1 = prov2.getDecryptionMaterials(ctx(eMat3_1, attr1));
    final DecryptionMaterials dMat3_2 = prov2.getDecryptionMaterials(ctx(eMat3_2, attr2));
    assertEquals(eMat3_1.getEncryptionKey(), dMat3_1.getDecryptionKey());
    assertEquals(eMat3_2.getEncryptionKey(), dMat3_2.getDecryptionKey());
    assertEquals(eMat3_1.getSigningKey(), dMat3_1.getVerificationKey());
    assertEquals(eMat3_2.getSigningKey(), dMat3_2.getVerificationKey());
    verifyNoMoreInteractions(client);
}