Merge pull request #513 from Parsely/testing/shared_consumer

stop producer tests from sharing a consumer instance

Author: emmettbutler

tests/pykafka/test_producer.py

@@ -7,6 +7,7 @@
 from uuid import uuid4
 
 from pykafka import KafkaClient
+from pykafka.common import OffsetType
 from pykafka.exceptions import MessageSizeTooLarge, ProducerQueueFullError
 from pykafka.partitioners import hashing_partitioner
 from pykafka.protocol import Message
@@ -26,27 +27,33 @@ def setUpClass(cls):
  cls.topic_name = b'test-data'
  cls.kafka.create_topic(cls.topic_name, 3, 2)
  cls.client = KafkaClient(cls.kafka.brokers, use_greenlets=cls.USE_GEVENT)
- cls.consumer = cls.client.topics[cls.topic_name].get_simple_consumer(
- consumer_timeout_ms=1000)
 
  @classmethod
  def tearDownClass(cls):
- cls.consumer.stop()
  stop_cluster(cls.kafka)
 
  def _get_producer(self, **kwargs):
  topic = self.client.topics[self.topic_name]
  return topic.get_producer(use_rdkafka=self.USE_RDKAFKA, **kwargs)
 
+ def _get_consumer(self):
+ return self.client.topics[self.topic_name].get_simple_consumer(
+ consumer_timeout_ms=1000,
+ auto_offset_reset=OffsetType.LATEST,
+ reset_offset_on_start=True,
+ )
+
  def test_produce(self):
  # unique bytes, just to be absolutely sure we're not fetching data
  # produced in a previous test
  payload = uuid4().bytes
 
+ consumer = self._get_consumer()
+
  prod = self._get_producer(sync=True, min_queued_messages=1)
  prod.produce(payload)
 
- message = self.consumer.consume()
+ message = consumer.consume()
  assert message.value == payload
 
  def test_sync_produce_raises(self):
@@ -60,41 +67,38 @@ def test_produce_hashing_partitioner(self):
  # produced in a previous test
  payload = uuid4().bytes
 
+ consumer = self._get_consumer()
+
  prod = self._get_producer(
  sync=True,
  min_queued_messages=1,
  partitioner=hashing_partitioner)
  prod.produce(payload, partition_key=b"dummy")
 
  # set a timeout so we don't wait forever if we break producer code
- message = self.consumer.consume()
+ message = consumer.consume()
  assert message.value == payload
 
  def test_async_produce(self):
  payload = uuid4().bytes
 
+ consumer = self._get_consumer()
+
  prod = self._get_producer(min_queued_messages=1, delivery_reports=True)
  prod.produce(payload)
 
  report = prod.get_delivery_report()
  self.assertEqual(report[0].value, payload)
  self.assertIsNone(report[1])
 
- message = self.consumer.consume()
+ message = consumer.consume()
  assert message.value == payload
 
  def test_recover_disconnected(self):
  """Test our retry-loop with a recoverable error"""
  payload = uuid4().bytes
  prod = self._get_producer(min_queued_messages=1, delivery_reports=True)
-
- # We must stop the consumer for this test, to ensure that it is the
- # producer that will encounter the disconnected brokers and initiate
- # a cluster update
- self.consumer.stop()
- for t in self.consumer._fetch_workers:
- t.join()
- part_offsets = self.consumer.held_offsets
+ consumer = self._get_consumer()
 
  for broker in self.client.brokers.values():
  broker._connection.disconnect()
@@ -103,56 +107,55 @@ def test_recover_disconnected(self):
  report = prod.get_delivery_report()
  self.assertIsNone(report[1])
 
- self.consumer.start()
- self.consumer.reset_offsets(
- # This is just a reset_offsets, but works around issue #216:
- [(self.consumer.partitions[pid], offset if offset != -1 else -2)
- for pid, offset in part_offsets.items()])
- message = self.consumer.consume()
+ message = consumer.consume()
  self.assertEqual(message.value, payload)
 
  def test_async_produce_context(self):
  """Ensure that the producer works as a context manager"""
  payload = uuid4().bytes
 
+ consumer = self._get_consumer()
  with self._get_producer(min_queued_messages=1) as producer:
  producer.produce(payload)
 
- message = self.consumer.consume()
+ message = consumer.consume()
  assert message.value == payload
 
  def test_async_produce_queue_full(self):
  """Ensure that the producer raises an error when its queue is full"""
+ consumer = self._get_consumer()
  with self._get_producer(block_on_queue_full=False,
  max_queued_messages=1,
  linger_ms=1000) as producer:
  with self.assertRaises(ProducerQueueFullError):
  while True:
  producer.produce(uuid4().bytes)
- while self.consumer.consume() is not None:
+ while consumer.consume() is not None:
  time.sleep(.05)
 
  def test_async_produce_lingers(self):
  """Ensure that the context manager waits for linger_ms milliseconds"""
  linger = 3
+ consumer = self._get_consumer()
  with self._get_producer(linger_ms=linger * 1000) as producer:
  start = time.time()
  producer.produce(uuid4().bytes)
  producer.produce(uuid4().bytes)
  self.assertTrue(int(time.time() - start) >= int(linger))
- self.consumer.consume()
- self.consumer.consume()
+ consumer.consume()
+ consumer.consume()
 
  def test_async_produce_thread_exception(self):
  """Ensure that an exception on a worker thread is raised to the main thread"""
+ consumer = self._get_consumer()
  with self.assertRaises(AttributeError):
  with self._get_producer(min_queued_messages=1) as producer:
  # get some dummy data into the queue that will cause a crash
  # when flushed:
  msg = Message("stuff", partition_id=0)
  del msg.value
  producer._produce(msg)
- while self.consumer.consume() is not None:
+ while consumer.consume() is not None:
  time.sleep(.05)
 
  def test_required_acks(self):
@@ -171,13 +174,14 @@ def test_required_acks(self):
 
  def test_null_payloads(self):
  """Test that None is accepted as a null payload"""
+ consumer = self._get_consumer()
  prod = self._get_producer(sync=True, min_queued_messages=1)
  prod.produce(None)
- self.assertIsNone(self.consumer.consume().value)
+ self.assertIsNone(consumer.consume().value)
  prod.produce(None, partition_key=b"whatever")
- self.assertIsNone(self.consumer.consume().value)
+ self.assertIsNone(consumer.consume().value)
  prod.produce(b"") # empty string should be distinguished from None
- self.assertEqual(b"", self.consumer.consume().value)
+ self.assertEqual(b"", consumer.consume().value)
 
  def test_owned_broker_flush_message_larger_then_max_request_size(self):
  """Test that producer batches messages into the batches no larger then
@@ -239,18 +243,19 @@ def test_async_produce_compression_large_message(self):
  # TODO: make payload size bigger once pypy snappy compression issue is
  # fixed
  large_payload = b''.join([uuid4().bytes for i in range(5)])
+ consumer = self._get_consumer()
 
  prod = self._get_producer(
-  compression=CompressionType.SNAPPY,
-  delivery_reports=True
-  )
+ compression=CompressionType.SNAPPY,
+ delivery_reports=True
+ )
  prod.produce(large_payload)
 
  report = prod.get_delivery_report()
  self.assertEqual(report[0].value, large_payload)
  self.assertIsNone(report[1])
 
- message = self.consumer.consume()
+ message = consumer.consume()
  assert message.value == large_payload
 
  for i in range(10):
@@ -259,6 +264,7 @@ def test_async_produce_compression_large_message(self):
  # use retry logic to loop over delivery reports and ensure we can
  # produce a group of large messages
  reports = []
+
  def ensure_all_messages_produced():
  report = prod.get_delivery_report()
  reports.append(report)
@@ -271,15 +277,17 @@ def ensure_all_messages_produced():
 
  # cleanup and consumer all messages
  msgs = []
+
  def ensure_all_messages_consumed():
- msg = self.consumer.consume()
+ msg = consumer.consume()
  if msg:
  msgs.append(msg)
  assert len(msgs) == 10
  retry(ensure_all_messages_consumed, retry_time=15)
 
  def test_async_produce_large_message(self):
 
+ consumer = self._get_consumer()
  large_payload = b''.join([uuid4().bytes for i in range(50000)])
  assert len(large_payload) / 1024 / 1024 < 1.0
 
@@ -290,7 +298,7 @@ def test_async_produce_large_message(self):
  self.assertEqual(report[0].value, large_payload)
  self.assertIsNone(report[1])
 
- message = self.consumer.consume()
+ message = consumer.consume()
  assert message.value == large_payload
 
  for i in range(10):
@@ -299,6 +307,7 @@ def test_async_produce_large_message(self):
  # use retry logic to loop over delivery reports and ensure we can
  # produce a group of large messages
  reports = []
+
  def ensure_all_messages_produced():
  report = prod.get_delivery_report()
  reports.append(report)
@@ -311,8 +320,9 @@ def ensure_all_messages_produced():
 
  # cleanup and consumer all messages
  msgs = []
+
  def ensure_all_messages_consumed():
- msg = self.consumer.consume()
+ msg = consumer.consume()
  if msg:
  msgs.append(msg)
  assert len(msgs) == 10