@Test
public void updatesMapMembersWhenContainerUpdated() {

    ShapeId stringId = ShapeId.from("ns.foo#String");
    ShapeId containerId = ShapeId.from("ns.foo#Container");
    ShapeId keyId = ShapeId.from("ns.foo#Container$key");
    ShapeId valueId = ShapeId.from("ns.foo#Container$value");
    StringShape target = StringShape.builder().id(stringId).build();
    MemberShape keyMember = MemberShape.builder().id(keyId).target(stringId).build();
    MemberShape valueMember = MemberShape.builder().id(valueId).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 newKey = MemberShape.builder()
            .id(keyId)
            .target(stringId)
            .addTrait(new SensitiveTrait())
            .build();
    MemberShape newValue = MemberShape.builder()
            .id(valueId)
            .target(stringId)
            .addTrait(new SensitiveTrait())
            .build();
    MapShape newMap = MapShape.builder().id(containerId).key(newKey).value(newValue).build();
    Model result = transformer.replaceShapes(model, Arrays.asList(newMap));
    assertThat(result.shapes().count(), Matchers.equalTo(4L));
    assertThat(result.getShape(keyId).get().getTrait(SensitiveTrait.class), Matchers.not(Optional.empty()));
    assertThat(result.getShape(valueId).get().getTrait(SensitiveTrait.class), Matchers.not(Optional.empty()));
    assertThat(result.getShape(containerId).get(), Matchers.is(newMap));
    assertThat(result.getShape(containerId).get().asMapShape().get().getKey(), Matchers.is(newKey));
    assertThat(result.getShape(containerId).get().asMapShape().get().getValue(), Matchers.is(newValue));
}