@Test
public void replacesMappedShapes() {

    ShapeId shapeId = ShapeId.from("ns.foo#id1");
    StringShape shape = StringShape.builder()
            .id(shapeId)
            .addTrait(new SensitiveTrait())
            .build();
    Model model = Model.builder().addShape(shape).build();
    ModelTransformer transformer = ModelTransformer.create();
    Model result = transformer.mapShapes(model, s -> Shape.shapeToBuilder(s).removeTrait("sensitive").build());
    assertThat(result.expectShape(shapeId).getId(), Matchers.equalTo(shapeId));
    assertThat(result.expectShape(shapeId).getTrait(SensitiveTrait.class), Matchers.is(Optional.empty()));
}