@Override
  public void close() throws Exception {
    Exception exception = null;
    for (CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>> receiver :
        ImmutableList.copyOf(receivers.values())) {
      // Cancel any observer waiting for the client to complete. If the receiver has already been
      // completed or cancelled, this call will be ignored.
      receiver.cancel(true);
      if (!receiver.isCompletedExceptionally()) {
        try {
          receiver.get().close();
        } catch (Exception e) {
          if (exception == null) {
            exception = e;
          } else {
            exception.addSuppressed(e);
          }
        }
      }
    }
    // Cancel any outbound calls and complete any inbound calls, as this multiplexer is hanging up
    outboundObserver.onError(
        Status.CANCELLED.withDescription("Multiplexer hanging up").asException());
    inboundObserver.onCompleted();
    if (exception != null) {
      throw exception;
    }
  }
@Test
  public void testClose() throws Exception {
    Collection<BeamFnApi.Elements> outboundValues = new ArrayList<>();
    Collection<Throwable> errorWasReturned = new ArrayList<>();
    AtomicBoolean wasClosed = new AtomicBoolean();
    final BeamFnDataGrpcMultiplexer multiplexer =
        new BeamFnDataGrpcMultiplexer(
            DESCRIPTOR,
            OutboundObserverFactory.clientDirect(),
            inboundObserver ->
                TestStreams.withOnNext(outboundValues::add)
                    .withOnError(errorWasReturned::add)
                    .build());
    multiplexer.registerConsumer(
        DATA_INSTRUCTION_ID,
        new CloseableFnDataReceiver<BeamFnApi.Elements>() {
          @Override
          public void flush() throws Exception {
            fail("Unexpected call");
          }

          @Override
          public void close() throws Exception {
            wasClosed.set(true);
          }

          @Override
          public void accept(BeamFnApi.Elements input) throws Exception {
            fail("Unexpected call");
          }
        });

    multiplexer.close();

    assertTrue(wasClosed.get());
    assertThat(
        Iterables.getOnlyElement(errorWasReturned).getMessage(),
        containsString("Multiplexer hanging up"));
  }