@Test
    public void testConcurrentPartitionAssignments() throws InterruptedException, ExecutionException {
        final List<TopicIdPartition> allPartitions = getIdPartitions("sample", 100);
        final Map<TopicPartition, Long> endOffsets = allPartitions.stream()
            .map(idp -> toRemoteLogPartition(partitioner.metadataPartition(idp)))
            .collect(Collectors.toMap(Function.identity(), e -> 0L, (a, b) -> b));
        consumer.updateEndOffsets(endOffsets);

        final AtomicBoolean isAllPartitionsAssigned = new AtomicBoolean(false);
        CountDownLatch latch = new CountDownLatch(1);
        Thread assignor = new Thread(() -> {
            int partitionsAssigned = 0;
            for (TopicIdPartition partition : allPartitions) {
                if (partitionsAssigned == 50) {
                    // Once half the topic partitions are assigned, wait for the consumer to catch up. This ensures
                    // that the consumer is already running when the rest of the partitions are assigned.
                    try {
                        latch.await(1, TimeUnit.MINUTES);
                    } catch (InterruptedException e) {
                        fail(e.getMessage());
                    }
                }
                consumerTask.addAssignmentsForPartitions(Collections.singleton(partition));
                partitionsAssigned++;
            }
            isAllPartitionsAssigned.set(true);
        });
        Runnable consumerRunnable = () -> {
            try {
                while (!isAllPartitionsAssigned.get()) {
                    consumerTask.maybeWaitForPartitionAssignments();
                    latch.countDown();
                }
            } catch (Exception e) {
                fail(e.getMessage());
            }
        };

        ExecutorService consumerExecutor = Executors.newSingleThreadExecutor();
        Future<?> future = consumerExecutor.submit(consumerRunnable);
        assignor.start();

        assignor.join();
        future.get();
    }
Tests that the consumer task shuts down gracefully when there were no assignments.
storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/ConsumerTaskTest.java