Follow up to PR 1782 -- fix tests (dpkp#1914)

dpkp · web-flow · commit 975087b4ae22 · 2019-09-30T11:18:46.000-07:00
diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py
@@ -268,6 +268,7 @@ def _retrieve_offsets(self, timestamps, timeout_ms=float("inf")):
 
         start_time = time.time()
         remaining_ms = timeout_ms
+        timestamps = copy.copy(timestamps)
         while remaining_ms > 0:
             if not timestamps:
                 return {}
@@ -294,7 +295,7 @@ def _retrieve_offsets(self, timestamps, timeout_ms=float("inf")):
                 if refresh_future.succeeded() and isinstance(future.exception, Errors.StaleMetadata):
                     log.debug("Stale metadata was raised, and we now have an updated metadata. Rechecking partition existance")
                     unknown_partition = future.exception.args[0]  # TopicPartition from StaleMetadata
-                    if not self._client.cluster.leader_for_partition(unknown_partition):
+                    if self._client.cluster.leader_for_partition(unknown_partition) is None:
                         log.debug("Removed partition %s from offsets retrieval" % (unknown_partition, ))
                         timestamps.pop(unknown_partition)
             else:
diff --git a/test/test_consumer_integration.py b/test/test_consumer_integration.py
@@ -766,12 +766,11 @@ def test_kafka_consumer_offsets_for_time_old(kafka_consumer, topic):
 @pytest.mark.skipif(env_kafka_version() < (0, 10, 1), reason="Requires KAFKA_VERSION >= 0.10.1")
 def test_kafka_consumer_offsets_for_times_errors(kafka_consumer_factory, topic):
     consumer = kafka_consumer_factory(fetch_max_wait_ms=200,
-                                    request_timeout_ms=500)
+                                      request_timeout_ms=500)
     tp = TopicPartition(topic, 0)
     bad_tp = TopicPartition(topic, 100)
 
     with pytest.raises(ValueError):
         consumer.offsets_for_times({tp: -1})
 
-    with pytest.raises(KafkaTimeoutError):
-        consumer.offsets_for_times({bad_tp: 0})
+    assert consumer.offsets_for_times({bad_tp: 0}) == {bad_tp: None}
