diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperationTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperationTest.java index b058c836abd4..a9951718b1c0 100644 --- a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperationTest.java +++ b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperationTest.java @@ -39,6 +39,7 @@ import java.util.HashMap; import java.util.Map; import java.util.UUID; +import java.util.function.Consumer; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.Mock; @@ -131,30 +132,22 @@ public void generateRequest_withIndex_throwsIllegalArgumentException() { public void generateRequest_nullValuesNotIgnoredByDefault() { FakeItemWithSort item = createUniqueFakeItemWithSort(); item.setOtherAttribute1("value-1"); - UpdateItemOperation updateItemOperation = UpdateItemOperation.create( - UpdateItemEnhancedRequest.builder(FakeItemWithSort.class).item(item).build()); - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build()); + UpdateItemOperation updateItemOperation = + UpdateItemOperation.create(requestFakeItemWithSort(item, b -> { })); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); + + String expectedUpdateExpression = "SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + + " REMOVE " + OTHER_ATTRIBUTE_2_NAME; Map expectedValues = new HashMap<>(); expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build()); Map expectedNames = new HashMap<>(); expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2"); - UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .expressionAttributeValues(expectedValues) - .expressionAttributeNames(expectedNames) - .key(expectedKey) - .returnValues(ReturnValue.ALL_NEW); - UpdateItemRequest expectedRequest = - baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + - " REMOVE " + OTHER_ATTRIBUTE_2_NAME) - .build(); - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest.Builder expectedRequestBuilder = ddbRequestBuilder(ddbKey(item.getId(), item.getSort())); + expectedRequestBuilder.expressionAttributeValues(expectedValues); + expectedRequestBuilder.expressionAttributeNames(expectedNames); + expectedRequestBuilder.updateExpression(expectedUpdateExpression); + UpdateItemRequest expectedRequest = expectedRequestBuilder.build(); assertThat(request, is(expectedRequest)); } @@ -164,34 +157,24 @@ public void generateRequest_withConditionExpression() { FakeItemWithSort item = createUniqueFakeItemWithSort(); item.setOtherAttribute1("value-1"); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) - .item(item) - .conditionExpression(CONDITION_EXPRESSION) - .build()); - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build()); + UpdateItemOperation.create(requestFakeItemWithSort(item, b -> b.conditionExpression(CONDITION_EXPRESSION))); + + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); + + + String expectedUpdateExpression = "SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + + " REMOVE " + OTHER_ATTRIBUTE_2_NAME; Map expectedValues = new HashMap<>(CONDITION_EXPRESSION.expressionValues()); expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build()); Map expectedNames = new HashMap<>(CONDITION_EXPRESSION.expressionNames()); expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2"); - UpdateItemRequest.Builder baseExpectedRequest = - UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .expressionAttributeValues(expectedValues) - .expressionAttributeNames(expectedNames) - .conditionExpression(CONDITION_EXPRESSION.expression()) - .key(expectedKey) - .returnValues(ReturnValue.ALL_NEW); - UpdateItemRequest expectedRequest = - baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + - " REMOVE " + OTHER_ATTRIBUTE_2_NAME) - .build(); - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest.Builder expectedRequestBuilder = ddbRequestBuilder(ddbKey(item.getId(), item.getSort())); + expectedRequestBuilder.expressionAttributeValues(expectedValues); + expectedRequestBuilder.expressionAttributeNames(expectedNames); + expectedRequestBuilder.updateExpression(expectedUpdateExpression); + expectedRequestBuilder.conditionExpression(CONDITION_EXPRESSION.expression()); + UpdateItemRequest expectedRequest = expectedRequestBuilder.build(); assertThat(request, is(expectedRequest)); } @@ -202,20 +185,16 @@ public void generateRequest_withMinimalConditionExpression() { item.setOtherAttribute1("value-1"); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) - .item(item) - .conditionExpression(MINIMAL_CONDITION_EXPRESSION) - .build()); + UpdateItemOperation.create(requestFakeItemWithSort(item, b -> b.conditionExpression(MINIMAL_CONDITION_EXPRESSION))); - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); Map expectedValues = new HashMap<>(); expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build()); Map expectedNames = new HashMap<>(); expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2"); + assertThat(request.conditionExpression(), is(MINIMAL_CONDITION_EXPRESSION.expression())); assertThat(request.expressionAttributeNames(), is(expectedNames)); assertThat(request.expressionAttributeValues(), is(expectedValues)); @@ -226,32 +205,22 @@ public void generateRequest_explicitlyUnsetIgnoreNulls() { FakeItemWithSort item = createUniqueFakeItemWithSort(); item.setOtherAttribute1("value-1"); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) - .item(item) - .ignoreNulls(false) - .build()); - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build()); + UpdateItemOperation.create(requestFakeItemWithSort(item, b-> b.ignoreNulls(false))); + + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); + + String expectedUpdateExpression = "SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + + " REMOVE " + OTHER_ATTRIBUTE_2_NAME; Map expectedValues = new HashMap<>(); expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build()); Map expectedNames = new HashMap<>(); expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2"); - UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .expressionAttributeValues(expectedValues) - .expressionAttributeNames(expectedNames) - .key(expectedKey) - .returnValues(ReturnValue.ALL_NEW); - UpdateItemRequest expectedRequest = - baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + - " REMOVE " + OTHER_ATTRIBUTE_2_NAME) - .build(); - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest.Builder expectedRequestBuilder = ddbRequestBuilder(ddbKey(item.getId(), item.getSort())); + expectedRequestBuilder.expressionAttributeValues(expectedValues); + expectedRequestBuilder.expressionAttributeNames(expectedNames); + expectedRequestBuilder.updateExpression(expectedUpdateExpression); + UpdateItemRequest expectedRequest = expectedRequestBuilder.build(); assertThat(request, is(expectedRequest)); } @@ -262,37 +231,27 @@ public void generateRequest_multipleSetters() { item.setOtherAttribute1("value-1"); item.setOtherAttribute2("value-2"); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) - .item(item) - .ignoreNulls(false) - .build()); - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build()); + UpdateItemOperation.create(requestFakeItemWithSort(item, b-> b.ignoreNulls(false))); + + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); + Map expectedValues = new HashMap<>(); expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build()); expectedValues.put(OTHER_ATTRIBUTE_2_VALUE, AttributeValue.builder().s("value-2").build()); Map expectedNames = new HashMap<>(); expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2"); - UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .expressionAttributeValues(expectedValues) - .expressionAttributeNames(expectedNames) - .key(expectedKey) - .returnValues(ReturnValue.ALL_NEW); - UpdateItemRequest expectedRequest1 = - baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + - ", " + OTHER_ATTRIBUTE_2_NAME + " = " + OTHER_ATTRIBUTE_2_VALUE) - .build(); - UpdateItemRequest expectedRequest2 = - baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_2_NAME + " = " + OTHER_ATTRIBUTE_2_VALUE + - ", " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE) - .build(); - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + + UpdateItemRequest.Builder expectedRequestBuilder = ddbRequestBuilder(ddbKey(item.getId(), item.getSort())); + expectedRequestBuilder.expressionAttributeValues(expectedValues); + expectedRequestBuilder.expressionAttributeNames(expectedNames); + + String updateExpression1 = "SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE + + ", " + OTHER_ATTRIBUTE_2_NAME + " = " + OTHER_ATTRIBUTE_2_VALUE; + UpdateItemRequest expectedRequest1 = expectedRequestBuilder.updateExpression(updateExpression1).build(); + String updateExpression2 = "SET " + OTHER_ATTRIBUTE_2_NAME + " = " + OTHER_ATTRIBUTE_2_VALUE + + ", " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE; + UpdateItemRequest expectedRequest2 = expectedRequestBuilder.updateExpression(updateExpression2).build(); assertThat(request, either(is(expectedRequest1)).or(is(expectedRequest2))); } @@ -301,31 +260,18 @@ public void generateRequest_multipleSetters() { public void generateRequest_multipleDeletes() { FakeItemWithSort item = createUniqueFakeItemWithSort(); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) - .item(item) - .ignoreNulls(false) - .build()); - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build()); + UpdateItemOperation.create(requestFakeItemWithSort(item, b-> b.ignoreNulls(false))); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); + Map expectedNames = new HashMap<>(); expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2"); - UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .expressionAttributeNames(expectedNames) - .key(expectedKey) - .returnValues(ReturnValue.ALL_NEW); - UpdateItemRequest expectedRequest1 = - baseExpectedRequest.updateExpression("REMOVE " + OTHER_ATTRIBUTE_1_NAME + ", " + OTHER_ATTRIBUTE_2_NAME) - .build(); - UpdateItemRequest expectedRequest2 = - baseExpectedRequest.updateExpression("REMOVE " + OTHER_ATTRIBUTE_2_NAME + ", " + OTHER_ATTRIBUTE_1_NAME) - .build(); - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest.Builder expectedRequestBuilder = ddbRequestBuilder(ddbKey(item.getId(), item.getSort())); + expectedRequestBuilder.expressionAttributeNames(expectedNames); + String updateExpression1 = "REMOVE " + OTHER_ATTRIBUTE_1_NAME + ", " + OTHER_ATTRIBUTE_2_NAME; + UpdateItemRequest expectedRequest1 = expectedRequestBuilder.updateExpression(updateExpression1).build(); + String updateExpression2 = "REMOVE " + OTHER_ATTRIBUTE_2_NAME + ", " + OTHER_ATTRIBUTE_1_NAME; + UpdateItemRequest expectedRequest2 = expectedRequestBuilder.updateExpression(updateExpression2).build(); assertThat(request,either(is(expectedRequest1)).or(is(expectedRequest2))); } @@ -335,29 +281,18 @@ public void generateRequest_canIgnoreNullValues() { FakeItemWithSort item = createUniqueFakeItemWithSort(); item.setOtherAttribute1("value-1"); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) - .item(item) - .ignoreNulls(true) - .build()); - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build()); - Map expectedValues = - singletonMap(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build()); - Map expectedNames = singletonMap(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); - UpdateItemRequest expectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE) - .expressionAttributeValues(expectedValues) - .expressionAttributeNames(expectedNames) - .key(expectedKey) - .returnValues(ReturnValue.ALL_NEW) - .build(); + UpdateItemOperation.create(requestFakeItemWithSort(item, b-> b.ignoreNulls(true))); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); + String expectedUpdateExpression = "SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE; + Map expectedValues = singletonMap(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build()); + Map expectedNames = singletonMap(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1"); - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest.Builder expectedRequestBuilder = ddbRequestBuilder(ddbKey(item.getId(), item.getSort())); + expectedRequestBuilder.expressionAttributeValues(expectedValues); + expectedRequestBuilder.expressionAttributeNames(expectedNames); + expectedRequestBuilder.updateExpression(expectedUpdateExpression); + UpdateItemRequest expectedRequest = expectedRequestBuilder.build(); assertThat(request, is(expectedRequest)); } @@ -366,23 +301,9 @@ public void generateRequest_canIgnoreNullValues() { public void generateRequest_keyOnlyItem() { FakeItemWithSort item = createUniqueFakeItemWithSort(); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) - .item(item) - .ignoreNulls(true) - .build()); - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build()); - UpdateItemRequest expectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(expectedKey) - .returnValues(ReturnValue.ALL_NEW) - .build(); - - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemOperation.create(requestFakeItemWithSort(item, b-> b.ignoreNulls(true))); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(), PRIMARY_CONTEXT, null); + UpdateItemRequest expectedRequest = ddbRequestBuilder(ddbKey(item.getId(), item.getSort())).build(); assertThat(request, is(expectedRequest)); } @@ -407,12 +328,7 @@ public void generateRequest_withExtension_modifiesKeyPortionOfItem() { mockDynamoDbEnhancedClientExtension); assertThat(request.key(), is(keyMap)); - verify(mockDynamoDbEnhancedClientExtension).beforeWrite(DefaultDynamoDbExtensionContext.builder() - .tableSchema(FakeItem.getTableSchema()) - .tableMetadata(FakeItem.getTableMetadata()) - .operationContext(PRIMARY_CONTEXT) - .operationName(OperationName.UPDATE_ITEM) - .items(baseMap).build()); + verify(mockDynamoDbEnhancedClientExtension).beforeWrite(extensionContext(baseMap, b -> b.operationName(OperationName.UPDATE_ITEM))); } @Test @@ -426,12 +342,8 @@ public void generateRequest_withExtension_modifiesUpdateExpression() { when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class))) .thenReturn(WriteModification.builder().transformedItem(fakeMap).build()); - UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(fakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(fakeItem, b -> b.ignoreNulls(true))); UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, @@ -443,25 +355,6 @@ public void generateRequest_withExtension_modifiesUpdateExpression() { assertThat(request.expressionAttributeNames(), hasEntry(SUBCLASS_ATTRIBUTE_NAME, "subclass_attribute")); } - @Test - public void transformResponse_mapsAttributesReturnedInResponse() { - FakeItem fakeItem1 = FakeItem.createUniqueFakeItem(); - FakeItem fakeItem2 = FakeItem.createUniqueFakeItem(); - Map fakeItem2Attributes = FakeItem.getTableSchema().itemToMap(fakeItem2, true); - - UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class).item(fakeItem1).build()); - - FakeItem result = updateItemOperation.transformResponse( - UpdateItemResponse.builder().attributes(fakeItem2Attributes).build(), - FakeItem.getTableSchema(), - PRIMARY_CONTEXT, - null) - .attributes(); - - assertThat(result, is(fakeItem2)); - } - @Test public void generateRequest_withExtensions_singleCondition() { FakeItem baseFakeItem = createUniqueFakeItem(); @@ -471,10 +364,7 @@ public void generateRequest_withExtensions_singleCondition() { when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class))) .thenReturn(WriteModification.builder().additionalConditionalExpression(condition).build()); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(baseFakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(baseFakeItem, b -> b.ignoreNulls(true))); UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, @@ -495,10 +385,7 @@ public void generateRequest_withExtensions_conflictingExpressionValue_throwsRunt .thenReturn(WriteModification.builder().additionalConditionalExpression(condition1).build()); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(baseFakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(baseFakeItem, b -> b.ignoreNulls(true))); try { updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, mockDynamoDbEnhancedClientExtension); @@ -519,10 +406,7 @@ public void generateRequest_withExtensions_conflictingExpressionName_throwsRunti .thenReturn(WriteModification.builder().additionalConditionalExpression(condition1).build()); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(baseFakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(baseFakeItem, b -> b.ignoreNulls(true))); try { updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, mockDynamoDbEnhancedClientExtension); @@ -544,10 +428,7 @@ public void generateRequest_withExtension_correctlyCoalescesIdenticalExpressionV .thenReturn(WriteModification.builder().additionalConditionalExpression(condition).build()); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(baseFakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(baseFakeItem, b -> b.ignoreNulls(true))); UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, @@ -566,10 +447,7 @@ public void generateRequest_withExtension_correctlyCoalescesIdenticalExpressionN .thenReturn(WriteModification.builder().additionalConditionalExpression(condition).build()); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(baseFakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(baseFakeItem, b -> b.ignoreNulls(true))); UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, @@ -585,10 +463,7 @@ public void generateRequest_withExtension_noModifications() { .thenReturn(WriteModification.builder().build()); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(baseFakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(baseFakeItem, b -> b.ignoreNulls(true))); UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, @@ -618,10 +493,7 @@ public void generateRequest_withExtension_conditionAndModification() { .build()); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(baseFakeItem) - .ignoreNulls(true) - .build()); + UpdateItemOperation.create(requestFakeItem(baseFakeItem, b -> b.ignoreNulls(true))); UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, @@ -641,25 +513,10 @@ public void generateRequest_withReturnConsumedCapacity_unknownValue_generatesCor String returnConsumedCapacity = UUID.randomUUID().toString(); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(item) - .ignoreNulls(true) - .returnConsumedCapacity(returnConsumedCapacity) - .build()); - - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - - UpdateItemRequest expectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(expectedKey).returnValues(ReturnValue.ALL_NEW) - .returnConsumedCapacity(returnConsumedCapacity) - .build(); - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), - PRIMARY_CONTEXT, - null); - + UpdateItemOperation.create(requestFakeItem(item, b -> b.ignoreNulls(true) + .returnConsumedCapacity(returnConsumedCapacity))); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, null); + UpdateItemRequest expectedRequest = ddbRequest(ddbKey(item.getId()), b -> b.returnConsumedCapacity(returnConsumedCapacity)); assertThat(request, is(expectedRequest)); } @@ -669,24 +526,11 @@ public void generateRequest_withReturnConsumedCapacity_knownValue_generatesCorre ReturnConsumedCapacity returnConsumedCapacity = ReturnConsumedCapacity.TOTAL; UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(item) - .ignoreNulls(true) - .returnConsumedCapacity(returnConsumedCapacity) - .build()); - - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); + UpdateItemOperation.create(requestFakeItem(item, b -> b.ignoreNulls(true) + .returnConsumedCapacity(returnConsumedCapacity))); - UpdateItemRequest expectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(expectedKey).returnValues(ReturnValue.ALL_NEW) - .returnConsumedCapacity(returnConsumedCapacity) - .build(); - - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, null); + UpdateItemRequest expectedRequest = ddbRequest(ddbKey(item.getId()), b -> b.returnConsumedCapacity(returnConsumedCapacity)); assertThat(request, is(expectedRequest)); } @@ -697,24 +541,11 @@ public void generateRequest_withReturnItemCollectionMetrics_unknownValue_generat String returnItemCollectionMetrics = UUID.randomUUID().toString(); UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(item) - .ignoreNulls(true) - .returnItemCollectionMetrics(returnItemCollectionMetrics) - .build()); - - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - - UpdateItemRequest expectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(expectedKey).returnValues(ReturnValue.ALL_NEW) - .returnItemCollectionMetrics(returnItemCollectionMetrics) - .build(); + UpdateItemOperation.create(requestFakeItem(item, b -> b.ignoreNulls(true) + .returnItemCollectionMetrics(returnItemCollectionMetrics))); - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, null); + UpdateItemRequest expectedRequest = ddbRequest(ddbKey(item.getId()), b -> b.returnItemCollectionMetrics(returnItemCollectionMetrics)); assertThat(request, is(expectedRequest)); } @@ -725,24 +556,11 @@ public void generateRequest_withReturnItemCollectionMetrics_knownValue_generates ReturnItemCollectionMetrics returnItemCollectionMetrics = ReturnItemCollectionMetrics.SIZE; UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class) - .item(item) - .ignoreNulls(true) - .returnItemCollectionMetrics(returnItemCollectionMetrics) - .build()); - - Map expectedKey = new HashMap<>(); - expectedKey.put("id", AttributeValue.builder().s(item.getId()).build()); - - UpdateItemRequest expectedRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(expectedKey).returnValues(ReturnValue.ALL_NEW) - .returnItemCollectionMetrics(returnItemCollectionMetrics) - .build(); + UpdateItemOperation.create(requestFakeItem(item, b -> b.ignoreNulls(true) + .returnItemCollectionMetrics(returnItemCollectionMetrics))); - UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), - PRIMARY_CONTEXT, - null); + UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(), PRIMARY_CONTEXT, null); + UpdateItemRequest expectedRequest = ddbRequest(ddbKey(item.getId()), b -> b.returnItemCollectionMetrics(returnItemCollectionMetrics)); assertThat(request, is(expectedRequest)); } @@ -754,72 +572,35 @@ public void transformResponse_withExtension_returnsCorrectTransformedItem() { Map baseFakeMap = FakeItem.getTableSchema().itemToMap(baseFakeItem, true); Map fakeMap = FakeItem.getTableSchema().itemToMap(fakeItem, true); - UpdateItemOperation 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) - .attributes(); + FakeItem resultItem = transformResponse(baseFakeItem); assertThat(resultItem, is(fakeItem)); - verify(mockDynamoDbEnhancedClientExtension).afterRead(DefaultDynamoDbExtensionContext.builder() - .tableMetadata(FakeItem.getTableMetadata()) - .tableSchema(FakeItem.getTableSchema()) - .operationContext(PRIMARY_CONTEXT) - .items(baseFakeMap).build()); + verify(mockDynamoDbEnhancedClientExtension).afterRead(extensionContext(baseFakeMap)); } @Test public void transformResponse_withNoOpExtension_returnsCorrectItem() { + when(mockDynamoDbEnhancedClientExtension.afterRead(any(DynamoDbExtensionContext.AfterRead.class))) + .thenReturn(ReadModification.builder().build()); + FakeItem baseFakeItem = createUniqueFakeItem(); Map baseFakeMap = FakeItem.getTableSchema().itemToMap(baseFakeItem, true); - UpdateItemOperation 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) - .attributes(); + FakeItem resultItem = transformResponse(baseFakeItem); assertThat(resultItem, is(baseFakeItem)); - verify(mockDynamoDbEnhancedClientExtension).afterRead(DefaultDynamoDbExtensionContext.builder() - .tableMetadata(FakeItem.getTableMetadata()) - .tableSchema(FakeItem.getTableSchema()) - .operationContext(PRIMARY_CONTEXT) - .items(baseFakeMap).build()); + verify(mockDynamoDbEnhancedClientExtension).afterRead(extensionContext(baseFakeMap)); } @Test(expected = IllegalStateException.class) public void transformResponse_afterReadThrowsException_throwsIllegalStateException() { - when(mockDynamoDbEnhancedClientExtension.afterRead(any(DynamoDbExtensionContext.AfterRead.class))).thenThrow(RuntimeException.class); - UpdateItemOperation updateItemOperation = - UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class).item(createUniqueFakeItem()).build()); + when(mockDynamoDbEnhancedClientExtension.afterRead(any(DynamoDbExtensionContext.AfterRead.class))) + .thenThrow(RuntimeException.class); - UpdateItemResponse response = - UpdateItemResponse.builder() - .attributes(FakeItem.getTableSchema().itemToMap(FakeItem.createUniqueFakeItem(), true)) - .build(); - - updateItemOperation.transformResponse(response, FakeItem.getTableSchema(), PRIMARY_CONTEXT, mockDynamoDbEnhancedClientExtension); + transformResponse(createUniqueFakeItem()); } @Test @@ -833,13 +614,12 @@ public void generateTransactWriteItem_basicRequest() { Map attributeValues = Collections.singletonMap("key", stringValue("value1")); Map attributeNames = Collections.singletonMap("key", "value2"); - UpdateItemRequest updateItemRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(fakeItemMap) - .updateExpression(updateExpression) - .expressionAttributeValues(attributeValues) - .expressionAttributeNames(attributeNames) - .build(); + UpdateItemRequest.Builder builder = ddbRequestBuilder(fakeItemMap); + builder.updateExpression(updateExpression); + builder.expressionAttributeValues(attributeValues); + builder.expressionAttributeNames(attributeNames); + UpdateItemRequest updateItemRequest = builder.build(); + doReturn(updateItemRequest).when(updateItemOperation).generateRequest(any(), any(), any()); TransactWriteItem actualResult = updateItemOperation.generateTransactWriteItem(FakeItem.getTableSchema(), @@ -871,14 +651,13 @@ public void generateTransactWriteItem_conditionalRequest() { Map attributeValues = Collections.singletonMap("key", stringValue("value1")); Map attributeNames = Collections.singletonMap("key", "value2"); - UpdateItemRequest updateItemRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(fakeItemMap) - .updateExpression(updateExpression) - .conditionExpression(conditionExpression) - .expressionAttributeValues(attributeValues) - .expressionAttributeNames(attributeNames) - .build(); + UpdateItemRequest.Builder builder = ddbRequestBuilder(fakeItemMap); + builder.updateExpression(updateExpression); + builder.expressionAttributeValues(attributeValues); + builder.expressionAttributeNames(attributeNames); + builder.conditionExpression(conditionExpression); + UpdateItemRequest updateItemRequest = builder.build(); + doReturn(updateItemRequest).when(updateItemOperation).generateRequest(any(), any(), any()); TransactWriteItem actualResult = updateItemOperation.generateTransactWriteItem(FakeItem.getTableSchema(), @@ -912,11 +691,7 @@ public void generateTransactWriteItem_returnValuesOnConditionCheckFailure_genera .build())); OperationContext context = DefaultOperationContext.create(TABLE_NAME, TableMetadata.primaryIndexName()); - UpdateItemRequest updateItemRequest = UpdateItemRequest.builder() - .tableName(TABLE_NAME) - .key(fakeItemMap) - .build(); - doReturn(updateItemRequest).when(updateItemOperation).generateRequest(any(), any(), any()); + doReturn(ddbRequest(fakeItemMap, b -> {})).when(updateItemOperation).generateRequest(any(), any(), any()); TransactWriteItem actualResult = updateItemOperation.generateTransactWriteItem(FakeItem.getTableSchema(), context, @@ -932,4 +707,72 @@ public void generateTransactWriteItem_returnValuesOnConditionCheckFailure_genera assertThat(actualResult, is(expectedResult)); verify(updateItemOperation).generateRequest(FakeItem.getTableSchema(), context, mockDynamoDbEnhancedClientExtension); } + + private Map ddbKey(String partitionKey) { + return singletonMap("id", AttributeValue.builder().s(partitionKey).build()); + } + + private Map ddbKey(String partitionKey, String sortKey) { + Map expectedKey = new HashMap<>(); + expectedKey.put("id", AttributeValue.builder().s(partitionKey).build()); + expectedKey.put("sort", AttributeValue.builder().s(sortKey).build()); + return expectedKey; + } + + private UpdateItemRequest ddbRequest(Map keys, Consumer modify) { + UpdateItemRequest.Builder builder = ddbBaseRequestBuilder(keys); + modify.accept(builder); + return builder.build(); + } + + private UpdateItemRequest.Builder ddbRequestBuilder(Map keys) { + return ddbBaseRequestBuilder(keys); + } + + private UpdateItemRequest.Builder ddbBaseRequestBuilder(Map keys) { + return UpdateItemRequest.builder().tableName(TABLE_NAME) + .key(keys) + .returnValues(ReturnValue.ALL_NEW); + } + + private UpdateItemEnhancedRequest requestFakeItem( + FakeItem item, Consumer> modify) { + UpdateItemEnhancedRequest.Builder builder = UpdateItemEnhancedRequest.builder(FakeItem.class).item(item); + modify.accept(builder); + return builder.build(); + } + + private UpdateItemEnhancedRequest requestFakeItemWithSort( + FakeItemWithSort item, Consumer> modify) { + UpdateItemEnhancedRequest.Builder builder = UpdateItemEnhancedRequest.builder(FakeItemWithSort.class) + .item(item); + modify.accept(builder); + return builder.build(); + } + + private DefaultDynamoDbExtensionContext extensionContext(Map fakeMap) { + return extensionContext(fakeMap, b -> {}); + } + + private DefaultDynamoDbExtensionContext extensionContext(Map fakeMap, + Consumer modify) { + DefaultDynamoDbExtensionContext.Builder builder = DefaultDynamoDbExtensionContext.builder() + .tableMetadata(FakeItem.getTableMetadata()) + .tableSchema(FakeItem.getTableSchema()) + .operationContext(PRIMARY_CONTEXT) + .items(fakeMap); + modify.accept(builder); + return builder.build(); + } + + private FakeItem transformResponse(FakeItem item) { + UpdateItemOperation updateItemOperation = + UpdateItemOperation.create(requestFakeItem(item, b -> b.ignoreNulls(true))); + + Map itemMap = FakeItem.getTableSchema().itemToMap(item, true); + return updateItemOperation.transformResponse(UpdateItemResponse.builder().attributes(itemMap).build(), + FakeItem.getTableSchema(), + PRIMARY_CONTEXT, + mockDynamoDbEnhancedClientExtension).attributes(); + } }