@Test
public void updatesTaggedUnionMembersWhenContainerUpdated() {

    ShapeId stringShapeId = ShapeId.from("ns.foo#String");
    StringShape string = StringShape.builder().id(stringShapeId).build();
    MemberShape member1 = MemberShape.builder().id("ns.foo#Shape$member1").target("ns.foo#String").build();
    MemberShape member2 = MemberShape.builder().id("ns.foo#Shape$member2").target("ns.foo#String").build();
    MemberShape member3 = MemberShape.builder().id("ns.foo#Shape$member3").target("ns.foo#String").build();
    UnionShape shape = UnionShape.builder()
            .id("ns.foo#Shape").addMember(member1).addMember(member2).addMember(member3).build();
    Model model = Model.builder().addShapes(string, shape, member1, member2, member3).build();
    // Add a trait to a replaced member3.
    MemberShape newMember3 = MemberShape.builder()
            .id("ns.foo#Shape$member3")
            .addTrait(new SensitiveTrait())
            .target("ns.foo#String")
            .build();
    // Replace the union with a shape that has the new member3.
    // Also remove member2.
    UnionShape other = UnionShape.builder()
            .id("ns.foo#Shape").addMember(member1).addMember(newMember3).build();
    Model result = ModelTransformer.create().replaceShapes(model, Collections.singleton(other));
    assertThat(result.getShape(member1.getId()), Matchers.equalTo(Optional.of(member1)));
    assertThat(result.getShape(member2.getId()), Matchers.is(Optional.empty()));
    assertThat(result.getShape(member3.getId()), Matchers.equalTo(Optional.of(newMember3)));
    assertThat(result.getShape(shape.getId()), Matchers.not(Optional.empty()));
}