public static void waitForCompletion(final KafkaStreams streams,
                                         final int expectedPartitions,
                                         final long timeoutMilliseconds) {
        final long start = System.currentTimeMillis();
        while (true) {
            int lagMetrics = 0;
            double totalLag = 0.0;
            for (final Metric metric : streams.metrics().values()) {
                if (metric.metricName().name().equals("records-lag")) {
                    if (!metric.metricName().tags().get("client-id").endsWith("restore-consumer")) {
                        lagMetrics++;
                        totalLag += ((Number) metric.metricValue()).doubleValue();
                    }
                }
            }
            if (lagMetrics >= expectedPartitions && totalLag == 0.0) {
                return;
            }
            if (System.currentTimeMillis() - start >= timeoutMilliseconds) {
                throw new RuntimeException(String.format(
                    "Timed out waiting for completion. lagMetrics=[%s/%s] totalLag=[%s]",
                    lagMetrics, expectedPartitions, totalLag
                ));
            }
        }
    }
Wait for streams to "finish", based on the consumer lag metric. Includes only the main consumer, for
completion of standbys as well see {@link #waitForStandbyCompletion}

Caveats:
- Inputs must be finite, fully loaded, and flushed before this method is called
- expectedPartitions is the total number of partitions to watch the lag on, including both input and internal.
  It's somewhat ok to get this wrong, as the main failure case would be an immediate return due to the clients
  not being initialized, which you can avoid with any non-zero value. But it's probably better to get it right ;)
streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java