@VisibleForTesting
  StreamingEngineConnectionState getCurrentConnections() {
    return connections.get();
  }
@Test
  public void testOnNewWorkerMetadata_correctlyRemovesStaleWindmillServers()
      throws InterruptedException {
    int metadataCount = 2;
    TestGetWorkBudgetDistributor getWorkBudgetDistributor =
        spy(new TestGetWorkBudgetDistributor(metadataCount));
    fanOutStreamingEngineWorkProvider =
        newStreamingEngineClient(
            GetWorkBudget.builder().setItems(1).setBytes(1).build(),
            getWorkBudgetDistributor,
            noOpProcessWorkItemFn());

    String workerToken = "workerToken1";
    String workerToken2 = "workerToken2";
    String workerToken3 = "workerToken3";

    WorkerMetadataResponse firstWorkerMetadata =
        WorkerMetadataResponse.newBuilder()
            .setMetadataVersion(1)
            .addWorkEndpoints(
                WorkerMetadataResponse.Endpoint.newBuilder()
                    .setBackendWorkerToken(workerToken)
                    .build())
            .addWorkEndpoints(
                WorkerMetadataResponse.Endpoint.newBuilder()
                    .setBackendWorkerToken(workerToken2)
                    .build())
            .putAllGlobalDataEndpoints(DEFAULT)
            .build();
    WorkerMetadataResponse secondWorkerMetadata =
        WorkerMetadataResponse.newBuilder()
            .setMetadataVersion(2)
            .addWorkEndpoints(
                WorkerMetadataResponse.Endpoint.newBuilder()
                    .setBackendWorkerToken(workerToken3)
                    .build())
            .putAllGlobalDataEndpoints(DEFAULT)
            .build();

    getWorkerMetadataReady.await();
    fakeGetWorkerMetadataStub.injectWorkerMetadata(firstWorkerMetadata);
    fakeGetWorkerMetadataStub.injectWorkerMetadata(secondWorkerMetadata);
    waitForWorkerMetadataToBeConsumed(getWorkBudgetDistributor);
    StreamingEngineConnectionState currentConnections =
        fanOutStreamingEngineWorkProvider.getCurrentConnections();
    assertEquals(1, currentConnections.windmillConnections().size());
    assertEquals(1, currentConnections.windmillStreams().size());
    Set<String> workerTokens =
        fanOutStreamingEngineWorkProvider.getCurrentConnections().windmillConnections().values()
            .stream()
            .map(WindmillConnection::backendWorkerToken)
            .collect(Collectors.toSet());

    assertFalse(workerTokens.contains(workerToken));
    assertFalse(workerTokens.contains(workerToken2));
  }