public static Pipeline updateTransform(
      String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder) {
    Components.Builder resultComponents = originalPipeline.getComponents().toBuilder();
    for (Map.Entry<String, PTransform> pt :
        originalPipeline.getComponents().getTransformsMap().entrySet()) {
      if (pt.getValue().getSpec() != null && urn.equals(pt.getValue().getSpec().getUrn())) {
        MessageWithComponents updated =
            compositeBuilder.getReplacement(pt.getKey(), originalPipeline.getComponents());
        if (updated == null) {
          continue;
        }

        checkArgument(
            updated.getPtransform().getOutputsMap().equals(pt.getValue().getOutputsMap()),
            "A %s must produce all of the outputs of the original %s",
            TransformReplacement.class.getSimpleName(),
            PTransform.class.getSimpleName());
        removeSubtransforms(pt.getValue(), resultComponents);
        resultComponents
            .mergeFrom(updated.getComponents())
            .putTransforms(pt.getKey(), updated.getPtransform());
      }
    }
    return originalPipeline.toBuilder().setComponents(resultComponents).build();
  }
@Test
  public void replaceExistingCompositeSucceeds() {
    Pipeline p =
        Pipeline.newBuilder()
            .addRootTransformIds("root")
            .setComponents(
                Components.newBuilder()
                    .putTransforms(
                        "root",
                        PTransform.newBuilder()
                            .addSubtransforms("sub_first")
                            .setSpec(FunctionSpec.newBuilder().setUrn("beam:composite"))
                            .build())
                    .putTransforms(
                        "sub_first",
                        PTransform.newBuilder()
                            .setSpec(FunctionSpec.newBuilder().setUrn("beam:inner"))
                            .build()))
            .build();

    Pipeline pipeline =
        ProtoOverrides.updateTransform(
            "beam:composite",
            p,
            new TestReplacer(
                PTransform.newBuilder()
                    .addSubtransforms("foo")
                    .addSubtransforms("bar")
                    .setSpec(
                        FunctionSpec.getDefaultInstance()
                            .newBuilderForType()
                            .setUrn("beam:composite"))
                    .build(),
                Components.getDefaultInstance()));
    assertThat(
        pipeline.getComponents().getTransformsOrThrow("root").getSpec().getUrn(),
        equalTo("beam:composite"));
    assertThat(
        pipeline.getComponents().getTransformsOrThrow("root").getSubtransformsList(),
        contains("foo", "bar"));
  }