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

2 files changed, 4 insertions, 4 deletions

diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py
index f781d4c..5434c36 100644
--- a/kafka/consumer/fetcher.py
+++ b/kafka/consumer/fetcher.py
@@ -268,6 +268,7 @@ class Fetcher(six.Iterator):
 
         start_time = time.time()
         remaining_ms = timeout_ms
+        timestamps = copy.copy(timestamps)
         while remaining_ms > 0:
             if not timestamps:
                 return {}
@@ -294,7 +295,7 @@ class Fetcher(six.Iterator):
                 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
index c7e2ebf..a2b8f70 100644
--- 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}