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

@Test
public void testUpdateStatus_emptyResponse() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);
  UpdateItemRequest updateItemRequest =
      getUpdateItemReq(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  when(ddbClient.updateItem(updateItemRequest))
      .thenReturn(CompletableFuture.completedFuture(UpdateItemResponse.builder().build()));

  CompletableFuture<Schedule> response =
      scheduleManager.updateStatus(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  Throwable exception = getException(response);
  assertTrue(exception instanceof IllegalStateException);
  assertEquals(exception.getMessage(), "Status update successful but status isn't returned.");

  verify(ddbClient, times(1)).updateItem(updateItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

@Test
public void testUpdateStatus_failure() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);
  UpdateItemRequest updateItemRequest =
      getUpdateItemReq(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  Exception exception = new Exception();
  CompletableFuture<UpdateItemResponse> response = new CompletableFuture<>();
  response.completeExceptionally(exception);
  when(ddbClient.updateItem(updateItemRequest)).thenReturn(response);

  CompletableFuture<Schedule> ret =
      scheduleManager.updateStatus(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  assertSame(getException(ret), exception);

  verify(ddbClient, times(1)).updateItem(updateItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

@Override
public T transformResponse(UpdateItemResponse response,
                           TableSchema<T> tableSchema,
                           OperationContext operationContext,
                           DynamoDbEnhancedClientExtension extension) {
    try {
        return readAndTransformSingleItem(response.attributes(), tableSchema, operationContext, extension);
    } catch (RuntimeException e) {
        // With a partial update it's possible to update the record into a state that the mapper can no longer
        // read or validate. This is more likely to happen with signed and encrypted records that undergo partial
        // updates (that practice is discouraged for this reason).
        throw new IllegalStateException("Unable to read the new item returned by UpdateItem after the update "
                                        + "occurred. Rollbacks are not supported by this operation, therefore the "
                                        + "record may no longer be readable using this model.", e);
    }
}
 

@Test
public void transformResponse_withNoOpExtension_returnsCorrectItem() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> baseFakeMap = FakeItem.getTableSchema().itemToMap(baseFakeItem, true);

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(baseFakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    when(mockDynamoDbEnhancedClientExtension.afterRead(any(DynamoDbExtensionContext.AfterRead.class))).thenReturn(
        ReadModification.builder().build());
    UpdateItemResponse response = UpdateItemResponse.builder()
                                                    .attributes(baseFakeMap)
                                                    .build();

    FakeItem resultItem = updateItemOperation.transformResponse(response, FakeItem.getTableSchema(),
                                                                PRIMARY_CONTEXT, mockDynamoDbEnhancedClientExtension);

    assertThat(resultItem, is(baseFakeItem));
    verify(mockDynamoDbEnhancedClientExtension).afterRead(DefaultDynamoDbExtensionContext.builder()
                                                                                         .tableMetadata(FakeItem.getTableMetadata())
                                                                                         .operationContext(PRIMARY_CONTEXT)
                                                                                         .items(baseFakeMap).build());
}
 

@Test
public void transformResponse_mapsAttributesReturnedInResponse() {
    FakeItem fakeItem1 = FakeItem.createUniqueFakeItem();
    FakeItem fakeItem2 = FakeItem.createUniqueFakeItem();
    Map<String, AttributeValue> fakeItem2Attributes = FakeItem.getTableSchema().itemToMap(fakeItem2, true);

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class).item(fakeItem1).build());

    FakeItem result = updateItemOperation.transformResponse(
        UpdateItemResponse.builder().attributes(fakeItem2Attributes).build(),
        FakeItem.getTableSchema(),
        PRIMARY_CONTEXT,
        null);

    assertThat(result, is(fakeItem2));
}
 

@Test
public void testUpdateStatus() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);
  UpdateItemRequest updateItemRequest =
      getUpdateItemReq(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  when(ddbClient.updateItem(updateItemRequest))
      .thenReturn(
          CompletableFuture.completedFuture(
              UpdateItemResponse.builder()
                  .attributes(
                      ImmutableMap.of(
                          Attribute.JOB_STATUS.columnName,
                          AttributeValue.builder().s(JobStatus.ACQUIRED.name()).build()))
                  .build()));

  CompletableFuture<Schedule> response =
      scheduleManager.updateStatus(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);
  Assert.assertEquals(response.get(1, TimeUnit.SECONDS), schedule.withStatus(JobStatus.ACQUIRED));

  verify(ddbClient, times(1)).updateItem(updateItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

@Test
public void getServiceCall_makesTheRightCallAndReturnsResponse() {
    FakeItem item = createUniqueFakeItem();
    UpdateItemOperation<FakeItem> updateItemOperation = UpdateItemOperation.create(
        UpdateItemEnhancedRequest.builder(FakeItem.class).item(item).build());
    UpdateItemRequest updateItemRequest = UpdateItemRequest.builder().tableName(TABLE_NAME).build();
    UpdateItemResponse expectedResponse = UpdateItemResponse.builder().build();
    when(mockDynamoDbClient.updateItem(any(UpdateItemRequest.class))).thenReturn(expectedResponse);

    UpdateItemResponse response = updateItemOperation.serviceCall(mockDynamoDbClient).apply(updateItemRequest);

    assertThat(response, sameInstance(expectedResponse));
    verify(mockDynamoDbClient).updateItem(updateItemRequest);
}
 

TestItem(TableSchema<?> schema,
                 Object v2Bean,
                 UpdateItemResponse v2UpdateItemResponse,

                 Object v1Bean,
                 UpdateItemResult v1UpdateItemResult) {
    this.schema = schema;
    this.v2Bean = v2Bean;
    this.v2UpdateItemResponse = v2UpdateItemResponse;

    this.v1Bean = v1Bean;
    this.v1UpdateItemResult = v1UpdateItemResult;
}
 

@Test(expected = IllegalStateException.class)
public void transformResponse_afterReadThrowsException_throwsIllegalStateException() {
    when(mockDynamoDbEnhancedClientExtension.afterRead(any(DynamoDbExtensionContext.AfterRead.class))).thenThrow(RuntimeException.class);
    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class).item(createUniqueFakeItem()).build());

    UpdateItemResponse response =
        UpdateItemResponse.builder()
                          .attributes(FakeItem.getTableSchema().itemToMap(FakeItem.createUniqueFakeItem(), true))
                          .build();

    updateItemOperation.transformResponse(response, FakeItem.getTableSchema(), PRIMARY_CONTEXT, mockDynamoDbEnhancedClientExtension);
}
 

@Test
public void transformResponse_withExtension_returnsCorrectTransformedItem() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    FakeItem fakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> baseFakeMap = FakeItem.getTableSchema().itemToMap(baseFakeItem, true);
    Map<String, AttributeValue> fakeMap = FakeItem.getTableSchema().itemToMap(fakeItem, true);

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(baseFakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    when(mockDynamoDbEnhancedClientExtension.afterRead(any(DynamoDbExtensionContext.AfterRead.class))).thenReturn(
        ReadModification.builder().transformedItem(fakeMap).build());
    UpdateItemResponse response = UpdateItemResponse.builder()
                                                    .attributes(baseFakeMap)
                                                    .build();

    FakeItem resultItem = updateItemOperation.transformResponse(response, FakeItem.getTableSchema(),
                                                                PRIMARY_CONTEXT,
                                                                mockDynamoDbEnhancedClientExtension);

    assertThat(resultItem, is(fakeItem));
    verify(mockDynamoDbEnhancedClientExtension).afterRead(DefaultDynamoDbExtensionContext.builder()
                                                                                         .tableMetadata(FakeItem.getTableMetadata())
                                                                                         .operationContext(PRIMARY_CONTEXT)
                                                                                         .items(baseFakeMap).build());
}
 

@Test
public void testUpdateStatus_unknownResponse() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);
  UpdateItemRequest updateItemRequest =
      getUpdateItemReq(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  when(ddbClient.updateItem(updateItemRequest))
      .thenReturn(
          CompletableFuture.completedFuture(
              UpdateItemResponse.builder()
                  .attributes(
                      ImmutableMap.of(
                          Attribute.JOB_STATUS.columnName,
                          AttributeValue.builder().s("magic").build()))
                  .build()));

  CompletableFuture<Schedule> response =
      scheduleManager.updateStatus(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  Throwable exception = getException(response);
  assertTrue(exception instanceof IllegalArgumentException);
  assertEquals(
      exception.getMessage(),
      "No enum constant com.airbnb.dynein.scheduler.Schedule.JobStatus.magic");

  verify(ddbClient, times(1)).updateItem(updateItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

@Test
public void updateEmptyStringWithCondition() {
    Map<String, AttributeValue> expectedItemMap = new HashMap<>();
    expectedItemMap.put("id", AttributeValue.builder().s("id123").build());
    expectedItemMap.put("s", EMPTY_STRING);
    TestBean testBean = new TestBean();
    testBean.setId("id123");
    testBean.setS("");

    UpdateItemResponse response = UpdateItemResponse.builder()
                                                    .attributes(expectedItemMap)
                                                    .build();
    when(mockDynamoDbClient.updateItem(any(UpdateItemRequest.class))).thenReturn(response);

    Expression conditionExpression = Expression.builder()
                                                 .expression("#attr = :val")
                                                 .expressionNames(singletonMap("#attr", "s"))
                                                 .expressionValues(singletonMap(":val", EMPTY_STRING))
                                                 .build();

    TestBean result = dynamoDbTable.updateItem(r -> r.item(testBean).conditionExpression(conditionExpression));

    Map<String, String> expectedExpressionAttributeNames = new HashMap<>();
    expectedExpressionAttributeNames.put("#AMZN_MAPPED_s", "s");
    expectedExpressionAttributeNames.put("#attr", "s");
    Map<String, AttributeValue> expectedExpressionAttributeValues = new HashMap<>();
    expectedExpressionAttributeValues.put(":AMZN_MAPPED_s", EMPTY_STRING);
    expectedExpressionAttributeValues.put(":val", EMPTY_STRING);
    Map<String, AttributeValue> expectedKeyMap = new HashMap<>();
    expectedKeyMap.put("id", AttributeValue.builder().s("id123").build());

    UpdateItemRequest expectedRequest =
        UpdateItemRequest.builder()
                         .tableName(TABLE_NAME)
                         .key(expectedKeyMap)
                         .returnValues(ReturnValue.ALL_NEW)
                         .updateExpression("SET #AMZN_MAPPED_s = :AMZN_MAPPED_s")
                         .conditionExpression("#attr = :val")
                         .expressionAttributeNames(expectedExpressionAttributeNames)
                         .expressionAttributeValues(expectedExpressionAttributeValues)
                         .build();

    verify(mockDynamoDbClient).updateItem(expectedRequest);
    assertThat(result.getId()).isEqualTo("id123");
    assertThat(result.getS()).isEmpty();
}
 

@Test
public void updateEmptyBytesWithCondition() {
    Map<String, AttributeValue> expectedItemMap = new HashMap<>();
    expectedItemMap.put("id", AttributeValue.builder().s("id123").build());
    expectedItemMap.put("b", EMPTY_BINARY);
    TestBean testBean = new TestBean();
    testBean.setId("id123");
    testBean.setB(EMPTY_BYTES);

    UpdateItemResponse response = UpdateItemResponse.builder()
                                                    .attributes(expectedItemMap)
                                                    .build();
    when(mockDynamoDbClient.updateItem(any(UpdateItemRequest.class))).thenReturn(response);

    Expression conditionExpression = Expression.builder()
                                                 .expression("#attr = :val")
                                                 .expressionNames(singletonMap("#attr", "b"))
                                                 .expressionValues(singletonMap(":val", EMPTY_BINARY))
                                                 .build();

    TestBean result = dynamoDbTable.updateItem(r -> r.item(testBean).conditionExpression(conditionExpression));

    Map<String, String> expectedExpressionAttributeNames = new HashMap<>();
    expectedExpressionAttributeNames.put("#AMZN_MAPPED_b", "b");
    expectedExpressionAttributeNames.put("#attr", "b");
    Map<String, AttributeValue> expectedExpressionAttributeValues = new HashMap<>();
    expectedExpressionAttributeValues.put(":AMZN_MAPPED_b", EMPTY_BINARY);
    expectedExpressionAttributeValues.put(":val", EMPTY_BINARY);
    Map<String, AttributeValue> expectedKeyMap = new HashMap<>();
    expectedKeyMap.put("id", AttributeValue.builder().s("id123").build());

    UpdateItemRequest expectedRequest =
        UpdateItemRequest.builder()
                         .tableName(TABLE_NAME)
                         .key(expectedKeyMap)
                         .returnValues(ReturnValue.ALL_NEW)
                         .updateExpression("SET #AMZN_MAPPED_b = :AMZN_MAPPED_b")
                         .conditionExpression("#attr = :val")
                         .expressionAttributeNames(expectedExpressionAttributeNames)
                         .expressionAttributeValues(expectedExpressionAttributeValues)
                         .build();

    verify(mockDynamoDbClient).updateItem(expectedRequest);
    assertThat(result.getId()).isEqualTo("id123");
    assertThat(result.getB()).isEqualTo(EMPTY_BYTES);
}
 

@Override
public Function<UpdateItemRequest, CompletableFuture<UpdateItemResponse>> asyncServiceCall(
    DynamoDbAsyncClient dynamoDbAsyncClient) {

    return dynamoDbAsyncClient::updateItem;
}
 

@Override
public Function<UpdateItemRequest, UpdateItemResponse> serviceCall(DynamoDbClient dynamoDbClient) {
    return dynamoDbClient::updateItem;
}
 

private static DynamoDbClient getV2Client(Blackhole bh, UpdateItemResponse updateItemResponse) {
    return new V2TestDynamoDbUpdateItemClient(bh, updateItemResponse);
}
 

public V2TestDynamoDbUpdateItemClient(Blackhole bh, UpdateItemResponse updateItemResponse) {
    super(bh);
    this.updateItemResponse = updateItemResponse;
}
 

@Override
public UpdateItemResponse updateItem(UpdateItemRequest updateItemRequest) {
    bh.consume(updateItemRequest);
    return this.updateItemResponse;
}