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 replacesOnlyMatching() {
    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:second"))
                            .build())
                    .putPcollections(
                        "intermediatePc",
                        PCollection.newBuilder().setUniqueName("intermediate").build())
                    .putCoders(
                        "coder",
                        Coder.newBuilder().setSpec(FunctionSpec.getDefaultInstance()).build()))
            .build();

    PTransform secondReplacement =
        PTransform.newBuilder()
            .addSubtransforms("second_sub")
            .setSpec(
                FunctionSpec.newBuilder()
                    .setUrn("beam:second:replacement")
                    .setPayload(
                        ByteString.copyFrom("foo-bar-baz".getBytes(StandardCharsets.UTF_8))))
            .build();
    WindowingStrategy introducedWS =
        WindowingStrategy.newBuilder()
            .setAccumulationMode(AccumulationMode.Enum.ACCUMULATING)
            .build();
    RunnerApi.Components extraComponents =
        Components.newBuilder()
            .putPcollections(
                "intermediatePc",
                PCollection.newBuilder().setUniqueName("intermediate_replacement").build())
            .putWindowingStrategies("new_ws", introducedWS)
            .putTransforms("second_sub", PTransform.getDefaultInstance())
            .build();

    Pipeline updated =
        ProtoOverrides.updateTransform(
            "beam:second", p, new TestReplacer(secondReplacement, extraComponents));
    PTransform updatedSecond = updated.getComponents().getTransformsOrThrow("second");

    assertThat(updatedSecond, equalTo(secondReplacement));
    assertThat(
        updated.getComponents().getWindowingStrategiesOrThrow("new_ws"), equalTo(introducedWS));
    assertThat(
        updated.getComponents().getTransformsOrThrow("second_sub"),
        equalTo(PTransform.getDefaultInstance()));

    // TODO: This might not be appropriate. Merging in the other direction might force that callers
    // are well behaved.
    assertThat(
        updated.getComponents().getPcollectionsOrThrow("intermediatePc").getUniqueName(),
        equalTo("intermediate_replacement"));

    // Assert that the untouched components are unchanged.
    assertThat(
        updated.getComponents().getTransformsOrThrow("first"),
        equalTo(p.getComponents().getTransformsOrThrow("first")));
    assertThat(
        updated.getComponents().getCodersOrThrow("coder"),
        equalTo(p.getComponents().getCodersOrThrow("coder")));
    assertThat(updated.getRootTransformIdsList(), equalTo(p.getRootTransformIdsList()));
  }