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 replacesMultiple() {
    RunnerApi.Pipeline p =
        Pipeline.newBuilder()
            .addAllRootTransformIds(ImmutableList.of("first", "second"))
            .setComponents(
                Components.newBuilder()
                    .putTransforms(
                        "first",
                        PTransform.newBuilder()
                            .setSpec(FunctionSpec.newBuilder().setUrn("beam:first"))
                            .build())
                    .putTransforms(
                        "second",
                        PTransform.newBuilder()
                            .setSpec(FunctionSpec.newBuilder().setUrn("beam:repeated"))
                            .build())
                    .putTransforms(
                        "third",
                        PTransform.newBuilder()
                            .setSpec(FunctionSpec.newBuilder().setUrn("beam:repeated"))
                            .build())
                    .putPcollections(
                        "intermediatePc",
                        PCollection.newBuilder().setUniqueName("intermediate").build())
                    .putCoders(
                        "coder",
                        Coder.newBuilder().setSpec(FunctionSpec.getDefaultInstance()).build()))
            .build();

    ByteString newPayload = ByteString.copyFrom("foo-bar-baz".getBytes(StandardCharsets.UTF_8));
    Pipeline updated =
        ProtoOverrides.updateTransform(
            "beam:repeated",
            p,
            (transformId, existingComponents) -> {
              String subtransform = String.format("%s_sub", transformId);
              return MessageWithComponents.newBuilder()
                  .setPtransform(
                      PTransform.newBuilder()
                          .setSpec(
                              FunctionSpec.newBuilder()
                                  .setUrn("beam:repeated:replacement")
                                  .setPayload(newPayload))
                          .addSubtransforms(subtransform))
                  .setComponents(
                      Components.newBuilder()
                          .putTransforms(
                              subtransform,
                              PTransform.newBuilder().setUniqueName(subtransform).build()))
                  .build();
            });
    PTransform updatedSecond = updated.getComponents().getTransformsOrThrow("second");
    PTransform updatedThird = updated.getComponents().getTransformsOrThrow("third");

    assertThat(updatedSecond, not(equalTo(p.getComponents().getTransformsOrThrow("second"))));
    assertThat(updatedThird, not(equalTo(p.getComponents().getTransformsOrThrow("third"))));
    assertThat(updatedSecond.getSubtransformsList(), contains("second_sub"));
    assertThat(updatedSecond.getSpec().getPayload(), equalTo(newPayload));
    assertThat(updatedThird.getSubtransformsList(), contains("third_sub"));
    assertThat(updatedThird.getSpec().getPayload(), equalTo(newPayload));

    assertThat(updated.getComponents().getTransformsMap(), hasKey("second_sub"));
    assertThat(updated.getComponents().getTransformsMap(), hasKey("third_sub"));
    assertThat(
        updated.getComponents().getTransformsOrThrow("second_sub").getUniqueName(),
        equalTo("second_sub"));
    assertThat(
        updated.getComponents().getTransformsOrThrow("third_sub").getUniqueName(),
        equalTo("third_sub"));
  }