apoorvmittal10 commented on code in PR #17870: URL: https://github.com/apache/kafka/pull/17870#discussion_r1913066020
########## core/src/test/java/kafka/server/share/DelayedShareFetchTest.java: ########## @@ -758,6 +765,285 @@ public void testTryCompleteWhenPartitionMaxBytesStrategyThrowsException() { Mockito.verify(delayedShareFetch, times(2)).releasePartitionLocks(any()); assertTrue(delayedShareFetch.lock().tryLock()); delayedShareFetch.lock().unlock(); + + assertTrue(future.isDone()); + assertFalse(future.isCompletedExceptionally()); + Map<TopicIdPartition, ShareFetchResponseData.PartitionData> partitionDataMap = future.join(); + assertEquals(1, partitionDataMap.size()); + assertTrue(partitionDataMap.containsKey(tp0)); + assertEquals("Exception thrown", partitionDataMap.get(tp0).errorMessage()); + } + + @Test + public void testPartitionMaxBytesFromUniformStrategyWhenAllPartitionsAreAcquirable() { + ReplicaManager replicaManager = mock(ReplicaManager.class); + String groupId = "grp"; + TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0)); + TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1)); + TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 2)); + TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 3)); + TopicIdPartition tp4 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 4)); + + SharePartition sp0 = mock(SharePartition.class); + SharePartition sp1 = mock(SharePartition.class); + SharePartition sp2 = mock(SharePartition.class); + SharePartition sp3 = mock(SharePartition.class); + SharePartition sp4 = mock(SharePartition.class); + + Map<TopicIdPartition, Integer> partitionMaxBytes = new HashMap<>(); + partitionMaxBytes.put(tp0, PARTITION_MAX_BYTES); + partitionMaxBytes.put(tp1, PARTITION_MAX_BYTES); + partitionMaxBytes.put(tp2, PARTITION_MAX_BYTES); + partitionMaxBytes.put(tp3, PARTITION_MAX_BYTES); + partitionMaxBytes.put(tp4, PARTITION_MAX_BYTES); + + when(sp0.maybeAcquireFetchLock()).thenReturn(true); + when(sp1.maybeAcquireFetchLock()).thenReturn(true); + when(sp2.maybeAcquireFetchLock()).thenReturn(true); + when(sp3.maybeAcquireFetchLock()).thenReturn(true); + when(sp4.maybeAcquireFetchLock()).thenReturn(true); + when(sp0.canAcquireRecords()).thenReturn(true); + when(sp1.canAcquireRecords()).thenReturn(true); + when(sp2.canAcquireRecords()).thenReturn(true); + when(sp3.canAcquireRecords()).thenReturn(true); + when(sp4.canAcquireRecords()).thenReturn(true); + + LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>(); + sharePartitions.put(tp0, sp0); + sharePartitions.put(tp1, sp1); + sharePartitions.put(tp2, sp2); + sharePartitions.put(tp3, sp3); + sharePartitions.put(tp4, sp4); + + ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, Uuid.randomUuid().toString(), + new CompletableFuture<>(), partitionMaxBytes, MAX_FETCH_RECORDS); + + when(sp0.acquire(anyString(), anyInt(), any(FetchPartitionData.class))).thenReturn( + ShareAcquiredRecords.fromAcquiredRecords(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1))); + when(sp1.acquire(anyString(), anyInt(), any(FetchPartitionData.class))).thenReturn( + ShareAcquiredRecords.fromAcquiredRecords(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1))); + when(sp2.acquire(anyString(), anyInt(), any(FetchPartitionData.class))).thenReturn( + ShareAcquiredRecords.fromAcquiredRecords(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1))); + when(sp3.acquire(anyString(), anyInt(), any(FetchPartitionData.class))).thenReturn( + ShareAcquiredRecords.fromAcquiredRecords(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1))); + when(sp4.acquire(anyString(), anyInt(), any(FetchPartitionData.class))).thenReturn( + ShareAcquiredRecords.fromAcquiredRecords(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1))); + + // All 5 partitions are acquirable. + doAnswer(invocation -> buildLogReadResult(sharePartitions.keySet())).when(replicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean()); Review Comment: Ahh, got it. Ignore the comment. -- 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