@Test
public void removesTraitsWhenDefinitionIsRemoved() {

    StringShape bazTrait = StringShape.builder()
            .id("ns.foo#baz")
            .addTrait(TraitDefinition.builder().build())
            .build();
    StringShape barTrait = StringShape.builder()
            .id("ns.foo#bar")
            .addTrait(TraitDefinition.builder().build())
            .build();
    ShapeId shapeId1 = ShapeId.from("ns.foo#id1");
    StringShape shape1 = StringShape.builder()
            .id(shapeId1)
            .addTrait(new DynamicTrait(ShapeId.from("foo.baz#foo"), Node.from(true)))
            .addTrait(new SensitiveTrait())
            .build();
    ShapeId shapeId2 = ShapeId.from("ns.foo#id2");
    StringShape shape2 = StringShape.builder()
            .id(shapeId2)
            .addTrait(new DynamicTrait(ShapeId.from("ns.foo#baz"), Node.from(true)))
            .addTrait(new DynamicTrait(ShapeId.from("ns.foo#bar"), Node.from(true)))
            .addTrait(new SensitiveTrait())
            .build();
    Model model = Model.builder()
            .addShapes(shape1, shape2, bazTrait, barTrait)
            .build();
    ModelTransformer transformer = ModelTransformer.create();
    Model result = transformer.filterShapes(model, shape -> !shape.getId().toString().equals("ns.foo#baz"));
    Map<Shape, TraitDefinition> definitions = new HashMap<>();
    for (Shape shape : result.getShapesWithTrait(TraitDefinition.class)) {
        definitions.put(shape, shape.expectTrait(TraitDefinition.class));
    }
    assertThat(definitions.size(), Matchers.is(1));
    assertThat(definitions, Matchers.hasKey(barTrait));
    assertThat(definitions, Matchers.hasValue(barTrait.getTrait(TraitDefinition.class).get()));
    assertThat(result.getShape(shapeId1).get().getTrait(SensitiveTrait.class), Matchers.not(Optional.empty()));
    assertThat(result.getShape(shapeId1).get().findTrait("ns.foo#baz"), Matchers.is(Optional.empty()));
    assertThat(result.getShape(shapeId2).get().getTrait(SensitiveTrait.class), Matchers.not(Optional.empty()));
    assertThat(result.getShape(shapeId2).get().findTrait("ns.foo#baz"), Matchers.is(Optional.empty()));
    assertThat(result.getShape(shapeId2).get().findTrait("ns.foo#bar"), Matchers.not(Optional.empty()));
}