Fix Map Iteration

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java

@@ -2006,15 +2006,12 @@ private void processTimestampSeeks() {
 				}
 			}
 			if (timestampSeeks != null) {
-				Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes = this.consumer
-						.offsetsForTimes(timestampSeeks);
-
-				for (TopicPartition tp : offsetsForTimes.keySet()) {
-					OffsetAndTimestamp ot = offsetsForTimes.get(tp);
+				Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes = this.consumer.offsetsForTimes(timestampSeeks);
+				offsetsForTimes.forEach((tp, ot) -> {
 					if (ot != null) {
 						this.consumer.seek(tp, ot.offset());
 					}
-				}
+				});
 			}
 		}
 

spring-kafka/src/test/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainerMockTests.java

@@ -372,10 +372,12 @@ void testAsyncTimestampSeeks() throws InterruptedException {
 				.collect(Collectors.toMap(tp -> tp, tp -> 200L)));
 		given(consumer.offsetsForTimes(Collections.singletonMap(tp0, 42L)))
 				.willReturn(Collections.singletonMap(tp0, new OffsetAndTimestamp(73L, 42L)));
+		Map<TopicPartition, OffsetAndTimestamp> map = new HashMap<>(assignments.stream()
+				.collect(Collectors.toMap(tp -> tp,
+						tp -> new OffsetAndTimestamp(tp.equals(tp0) ? 73L : 92L, 43L))));
+		map.put(new TopicPartition("foo", 5), null);
 		given(consumer.offsetsForTimes(any()))
-				.willReturn(assignments.stream()
-						.collect(Collectors.toMap(tp -> tp,
-								tp -> new OffsetAndTimestamp(tp.equals(tp0) ? 73L : 92L, 43L))));
+				.willReturn(map);
 		given(consumerFactory.createConsumer("grp", "", "-0", KafkaTestUtils.defaultPropertyOverrides()))
 				.willReturn(consumer);
 		ContainerProperties containerProperties = new ContainerProperties("foo");