@Test
public void updatesMapShapeWhenKeyOrValueIsUpdated() {

    ShapeId stringId = ShapeId.from("ns.foo#String");
    ShapeId containerId = ShapeId.from("ns.foo#Container");
    ShapeId keyMemberId = ShapeId.from("ns.foo#Container$key");
    ShapeId valueMemberId = ShapeId.from("ns.foo#Container$value");
    StringShape target = StringShape.builder().id(stringId).build();
    MemberShape keyMember = MemberShape.builder().id(keyMemberId).target(stringId).build();
    MemberShape valueMember = MemberShape.builder().id(valueMemberId).target(stringId).build();
    MapShape container = MapShape.builder().id(containerId).key(keyMember).value(valueMember).build();
    Model model = Model.builder()
            .addShapes(target, keyMember, valueMember, container)
            .build();
    ModelTransformer transformer = ModelTransformer.create();
    MemberShape newKeyMember = MemberShape.builder()
            .id(keyMemberId)
            .target(stringId)
            .addTrait(new SensitiveTrait())
            .build();
    Model resultWithNewKey = transformer.replaceShapes(model, Arrays.asList(newKeyMember));
    MemberShape newValueMember = MemberShape.builder()
            .id(valueMemberId)
            .target(stringId)
            .addTrait(new SensitiveTrait())
            .build();
    Model resultWithNewValue = transformer.replaceShapes(model, Arrays.asList(newValueMember));
    assertThat(resultWithNewKey.getShape(keyMemberId).get().getTrait(SensitiveTrait.class),
                      Matchers.not(Optional.empty()));
    assertThat(resultWithNewKey.getShape(containerId).get().asMapShape().get().getKey(),
                      Matchers.is(newKeyMember));
    assertThat(resultWithNewValue.getShape(valueMemberId).get().getTrait(SensitiveTrait.class),
                      Matchers.not(Optional.empty()));
    assertThat(resultWithNewValue.getShape(containerId).get().asMapShape().get().getValue(),
                      Matchers.is(newValueMember));
}