Repository: apache/kafka
File: streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStreamsRebalanceListenerTest.java (Java)

Code before:
        assertDoesNotThrow(() -> defaultStreamsRebalanceListener.onTasksRevoked(
            Set.of(new StreamsRebalanceData.TaskId("1", 0))
        ));

        final InOrder inOrder = inOrder(taskManager, streamThread);
        inOrder.verify(taskManager).handleRevocation(
            Set.of(new TopicPartition("source1", 0), new TopicPartition("repartition1", 0))
        );
        inOrder.verify(streamThread).state();
        if (state != StreamThread.State.PENDING_SHUTDOWN) {
            inOrder.verify(streamThread).setState(StreamThread.State.PARTITIONS_REVOKED);
        } else {
            inOrder.verify(streamThread, never()).setState(StreamThread.State.PARTITIONS_REVOKED);
        }
    }

    @Test
    void testOnTasksRevokedWithException() {
        final Exception exception = new RuntimeException("sample exception");
        doThrow(exception).when(taskManager).handleRevocation(any());

        createRebalanceListenerWithRebalanceData(new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Map.of(), Map.of()));

        final Exception actualException = assertThrows(RuntimeException.class, () -> defaultStreamsRebalanceListener.onTasksRevoked(Set.of()));

        assertEquals(actualException,  exception);
        verify(taskManager).handleRevocation(any());
        verify(streamThread, never()).setState(any());
    }

    @Test
    void testOnTasksAssigned() {
        final StreamsRebalanceData streamsRebalanceData = mock(StreamsRebalanceData.class);
        when(streamsRebalanceData.subtopologies()).thenReturn(Map.of(
            "1",
            new StreamsRebalanceData.Subtopology(
                Set.of("source1"),
                Set.of(),
                Map.of("repartition1", new StreamsRebalanceData.TopicInfo(Optional.of(1), Optional.of((short) 1), Map.of())),
                Map.of(),
                Set.of()
            ),
            "2",
            new StreamsRebalanceData.Subtopology(
                Set.of("source2"),
                Set.of(),
                Map.of("repartition2", new StreamsRebalanceData.TopicInfo(Optional.of(1), Optional.of((short) 1), Map.of())),
                Map.of(),
                Set.of()
            ),
            "3",

Review comment (suggestion):
[nitpick] Extra space before 'exception' parameter. Should be 'assertEquals(actualException, exception);' for consistent formatting.

Code after:
        assertDoesNotThrow(() -> defaultStreamsRebalanceListener.onTasksRevoked(
            Set.of(new StreamsRebalanceData.TaskId("1", 0))
        ));

        final InOrder inOrder = inOrder(taskManager, streamThread);
        inOrder.verify(taskManager).handleRevocation(
            Set.of(new TopicPartition("source1", 0), new TopicPartition("repartition1", 0))
        );
        inOrder.verify(streamThread).state();
        if (state != StreamThread.State.PENDING_SHUTDOWN) {
            inOrder.verify(streamThread).setState(StreamThread.State.PARTITIONS_REVOKED);
        } else {
            inOrder.verify(streamThread, never()).setState(StreamThread.State.PARTITIONS_REVOKED);
        }
    }

    @Test
    void testOnTasksRevokedWithException() {
        final Exception exception = new RuntimeException("sample exception");
        doThrow(exception).when(taskManager).handleRevocation(any());

        createRebalanceListenerWithRebalanceData(new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Map.of(), Map.of()));

        final Exception actualException = assertThrows(RuntimeException.class, () -> defaultStreamsRebalanceListener.onTasksRevoked(Set.of()));

        assertEquals(actualException, exception);
        verify(taskManager).handleRevocation(any());
        verify(streamThread, never()).setState(any());
    }

    @Test
    void testOnTasksAssigned() {
        final StreamsRebalanceData streamsRebalanceData = mock(StreamsRebalanceData.class);
        when(streamsRebalanceData.subtopologies()).thenReturn(Map.of(
            "1",
            new StreamsRebalanceData.Subtopology(
                Set.of("source1"),
                Set.of(),
                Map.of("repartition1", new StreamsRebalanceData.TopicInfo(Optional.of(1), Optional.of((short) 1), Map.of())),
                Map.of(),
                Set.of()
            ),
            "2",
            new StreamsRebalanceData.Subtopology(
                Set.of("source2"),
                Set.of(),
                Map.of("repartition2", new StreamsRebalanceData.TopicInfo(Optional.of(1), Optional.of((short) 1), Map.of())),
                Map.of(),
                Set.of()
            ),
            "3",