brenden20 commented on code in PR #16200:
URL: https://github.com/apache/kafka/pull/16200#discussion_r1669143238


##########
clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java:
##########
@@ -106,72 +102,90 @@ public class HeartbeatRequestManagerTest {
     private MembershipManager membershipManager;
     private HeartbeatRequestManager.HeartbeatRequestState 
heartbeatRequestState;
     private HeartbeatRequestManager.HeartbeatState heartbeatState;
-    private final String memberId = "member-id";
-    private final int memberEpoch = 1;
     private BackgroundEventHandler backgroundEventHandler;
-    private Metrics metrics;
+    private LogContext logContext;
 
     @BeforeEach
     public void setUp() {
-        setUp(ConsumerTestBuilder.createDefaultGroupInformation());
-    }
+        this.time = new MockTime();
+        this.logContext = new LogContext();
+        this.pollTimer = mock(Timer.class);
+        this.coordinatorRequestManager = mock(CoordinatorRequestManager.class);
+        this.heartbeatState = mock(HeartbeatState.class);
+        this.backgroundEventHandler = mock(BackgroundEventHandler.class);
+        this.subscriptions = mock(SubscriptionState.class);
+        this.membershipManager = mock(MembershipManagerImpl.class);
+        this.metadata = mock(ConsumerMetadata.class);
+        Metrics metrics = new Metrics(time);
+        ConsumerConfig config = mock(ConsumerConfig.class);
+
+        this.heartbeatRequestState = spy(new HeartbeatRequestState(
+                logContext,
+                time,
+                DEFAULT_HEARTBEAT_INTERVAL_MS,
+                DEFAULT_RETRY_BACKOFF_MS,
+                DEFAULT_RETRY_BACKOFF_MAX_MS,
+                DEFAULT_HEARTBEAT_JITTER_MS));
 
-    private void setUp(Optional<ConsumerTestBuilder.GroupInformation> 
groupInfo) {
-        testBuilder = new ConsumerTestBuilder(groupInfo, true, false);
-        time = testBuilder.time;
-        coordinatorRequestManager = 
testBuilder.coordinatorRequestManager.orElseThrow(IllegalStateException::new);
-        heartbeatRequestManager = 
testBuilder.heartbeatRequestManager.orElseThrow(IllegalStateException::new);
-        heartbeatRequestState = 
testBuilder.heartbeatRequestState.orElseThrow(IllegalStateException::new);
-        heartbeatState = 
testBuilder.heartbeatState.orElseThrow(IllegalStateException::new);
-        backgroundEventHandler = testBuilder.backgroundEventHandler;
-        subscriptions = testBuilder.subscriptions;
-        membershipManager = 
testBuilder.membershipManager.orElseThrow(IllegalStateException::new);
-        metadata = testBuilder.metadata;
-        metrics = new Metrics(time);
+        this.heartbeatRequestManager = new HeartbeatRequestManager(
+                logContext,
+                pollTimer,
+                config,
+                coordinatorRequestManager,
+                membershipManager,
+                heartbeatState,
+                heartbeatRequestState,
+                backgroundEventHandler,
+                metrics);
 
         
when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(new 
Node(1, "localhost", 9999)));
+        Map<Uuid, SortedSet<Integer>> map = new HashMap<>();
+        LocalAssignment local = new LocalAssignment(0, map);
+        when(membershipManager.currentAssignment()).thenReturn(local);
     }
 
-    private void resetWithZeroHeartbeatInterval(Optional<String> 
groupInstanceId) {
-        cleanup();
-
-        ConsumerTestBuilder.GroupInformation gi = new 
ConsumerTestBuilder.GroupInformation(
-                DEFAULT_GROUP_ID,
-                groupInstanceId,
+    private void createHeartbeatStateWith0HeartbeatInterval() {
+        this.heartbeatRequestState = spy(new HeartbeatRequestState(
+                logContext,
+                time,
                 0,
-                0.0,
-                Optional.of(DEFAULT_REMOTE_ASSIGNOR)
-        );
-
-        setUp(Optional.of(gi));
-    }
+                DEFAULT_RETRY_BACKOFF_MS,
+                DEFAULT_RETRY_BACKOFF_MAX_MS,
+                DEFAULT_HEARTBEAT_JITTER_MS));
 
-    @AfterEach
-    public void cleanup() {
-        if (testBuilder != null) {
-            testBuilder.close();
-        }
+        heartbeatRequestManager = createHeartbeatRequestManager(
+                coordinatorRequestManager,
+                membershipManager,
+                heartbeatState,
+                heartbeatRequestState,
+                backgroundEventHandler);
     }
 
     @Test
     public void testHeartbeatOnStartup() {
+        
when(coordinatorRequestManager.coordinator()).thenReturn(Optional.empty());
         NetworkClientDelegate.PollResult result = 
heartbeatRequestManager.poll(time.milliseconds());
         assertEquals(0, result.unsentRequests.size());
 
-        resetWithZeroHeartbeatInterval(Optional.empty());
-        mockStableMember();
+        time.sleep(DEFAULT_HEARTBEAT_INTERVAL_MS);
         assertEquals(0, 
heartbeatRequestManager.maximumTimeToWait(time.milliseconds()));
+        
when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mock(Node.class)));
         result = heartbeatRequestManager.poll(time.milliseconds());
         assertEquals(1, result.unsentRequests.size());
 
         // Ensure we do not resend the request without the first request being 
completed
+        
when(coordinatorRequestManager.coordinator()).thenReturn(Optional.empty());

Review Comment:
   Similar to above, I have made the change. Let me know if it looks good



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscr...@kafka.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

Reply via email to