<s> from __future__ import division <EOL> import platform <EOL> import pytest <EOL> import time <EOL> import unittest2 <EOL> from uuid import uuid4 <EOL> from pykafka import KafkaClient <EOL> from pykafka . common import OffsetType <EOL> from pykafka . exceptions import MessageSizeTooLarge , ProducerQueueFullError <EOL> from pykafka . partitioners import hashing_partitioner <EOL> from pykafka . protocol import Message <EOL> from pykafka . test . utils import get_cluster , stop_cluster , retry <EOL> from pykafka . common import CompressionType <EOL> from pykafka . producer import OwnedBroker <EOL> class ProducerIntegrationTests ( unittest2 . TestCase ) : <EOL> maxDiff = None <EOL> USE_RDKAFKA = False <EOL> USE_GEVENT = False <EOL> @ classmethod <EOL> def setUpClass ( cls ) : <EOL> cls . kafka = get_cluster ( ) <EOL> cls . topic_name = b'<STR_LIT>' <EOL> cls . kafka . create_topic ( cls . topic_name , <NUM_LIT:3> , <NUM_LIT:2> ) <EOL> cls . client = KafkaClient ( cls . kafka . brokers , use_greenlets = cls . USE_GEVENT ) <EOL> @ classmethod <EOL> def tearDownClass ( cls ) : <EOL> stop_cluster ( cls . kafka ) <EOL> def _get_producer ( self , ** kwargs ) : <EOL> topic = self . client . topics [ self . topic_name ] <EOL> return topic . get_producer ( use_rdkafka = self . USE_RDKAFKA , ** kwargs ) <EOL> def _get_consumer ( self ) : <EOL> return self . client . topics [ self . topic_name ] . get_simple_consumer ( <EOL> consumer_timeout_ms = <NUM_LIT:1000> , <EOL> auto_offset_reset = OffsetType . LATEST , <EOL> reset_offset_on_start = True , <EOL> ) <EOL> def test_produce ( self ) : <EOL> payload = uuid4 ( ) . bytes <EOL> consumer = self . _get_consumer ( ) <EOL> prod = self . _get_producer ( sync = True , min_queued_messages = <NUM_LIT:1> ) <EOL> prod . produce ( payload ) <EOL> message = consumer . consume ( ) <EOL> assert message . value == payload <EOL> def test_sync_produce_raises ( self ) : <EOL> """<STR_LIT>""" <EOL> with self . _get_producer ( sync = True , min_queued_messages = <NUM_LIT:1> ) as prod : <EOL> with self . assertRaises ( MessageSizeTooLarge ) : <EOL> prod . produce ( <NUM_LIT:10> ** <NUM_LIT:7> * b"<STR_LIT:U+0020>" ) <EOL> def test_produce_hashing_partitioner ( self ) : <EOL> payload = uuid4 ( ) . bytes <EOL> consumer = self . _get_consumer ( ) <EOL> prod = self . _get_producer ( <EOL> sync = True , <EOL> min_queued_messages = <NUM_LIT:1> , <EOL> partitioner = hashing_partitioner ) <EOL> prod . produce ( payload , partition_key = b"<STR_LIT>" ) <EOL> message = consumer . consume ( ) <EOL> assert message . value == payload <EOL> def test_async_produce ( self ) : <EOL> payload = uuid4 ( ) . bytes <EOL> consumer = self . _get_consumer ( ) <EOL> prod = self . _get_producer ( min_queued_messages = <NUM_LIT:1> , delivery_reports = True ) <EOL> prod . produce ( payload ) <EOL> report = prod . get_delivery_report ( ) <EOL> self . assertEqual ( report [ <NUM_LIT:0> ] . value , payload ) <EOL> self . assertIsNone ( report [ <NUM_LIT:1> ] ) <EOL> message = consumer . consume ( ) <EOL> assert message . value == payload <EOL> def test_recover_disconnected ( self ) : <EOL> """<STR_LIT>""" <EOL> payload = uuid4 ( ) . bytes <EOL> prod = self . _get_producer ( min_queued_messages = <NUM_LIT:1> , delivery_reports = True ) <EOL> consumer = self . _get_consumer ( ) <EOL> for broker in self . client . brokers . values ( ) : <EOL> broker . _connection . disconnect ( ) <EOL> prod . produce ( payload ) <EOL> report = prod . get_delivery_report ( ) <EOL> self . assertIsNone ( report [ <NUM_LIT:1> ] ) <EOL> message = consumer . consume ( ) <EOL> self . assertEqual ( message . value , payload ) <EOL> def test_async_produce_context ( self ) : <EOL> """<STR_LIT>""" <EOL> payload = uuid4 ( ) . bytes <EOL> consumer = self . _get_consumer ( ) <EOL> with self . _get_producer ( min_queued_messages = <NUM_LIT:1> ) as producer : <EOL> producer . produce ( payload ) <EOL> message = consumer . consume ( ) <EOL> assert message . value == payload <EOL> def test_async_produce_queue_full ( self ) : <EOL> """<STR_LIT>""" <EOL> consumer = self . _get_consumer ( ) <EOL> with self . _get_producer ( block_on_queue_full = False , <EOL> max_queued_messages = <NUM_LIT:1> , <EOL> linger_ms = <NUM_LIT:1000> ) as producer : <EOL> with self . assertRaises ( ProducerQueueFullError ) : <EOL> while True : <EOL> producer . produce ( uuid4 ( ) . bytes ) <EOL> while consumer . consume ( ) is not None : <EOL> time . sleep ( <NUM_LIT> ) <EOL> def test_async_produce_lingers ( self ) : <EOL> """<STR_LIT>""" <EOL> linger = <NUM_LIT:3> <EOL> consumer = self . _get_consumer ( ) <EOL> with self . _get_producer ( linger_ms = linger * <NUM_LIT:1000> ) as producer : <EOL> start = time . time ( ) <EOL> producer . produce ( uuid4 ( ) . bytes ) <EOL> producer . produce ( uuid4 ( ) . bytes ) <EOL> self . assertTrue ( int ( time . time ( ) - start ) >= int ( linger ) ) <EOL> consumer . consume ( ) <EOL> consumer . consume ( ) <EOL> def test_async_produce_thread_exception ( self ) : <EOL> """<STR_LIT>""" <EOL> consumer = self . _get_consumer ( ) <EOL> with self . assertRaises ( AttributeError ) : <EOL> with self . _get_producer ( min_queued_messages = <NUM_LIT:1> ) as producer : <EOL> msg = Message ( "<STR_LIT>" , partition_id = <NUM_LIT:0> ) <EOL> del msg . value <EOL> producer . _produce ( msg ) <EOL> while consumer . consume ( ) is not None : <EOL> time . sleep ( <NUM_LIT> ) <EOL> def test_required_acks ( self ) : <EOL> """<STR_LIT>""" <EOL> kwargs = dict ( linger_ms = <NUM_LIT:1> , sync = True , required_acks = <NUM_LIT:0> ) <EOL> prod = self . _get_producer ( ** kwargs ) <EOL> prod . produce ( uuid4 ( ) . bytes ) <EOL> kwargs [ "<STR_LIT>" ] = - <NUM_LIT:1> <EOL> prod = self . _get_producer ( ** kwargs ) <EOL> prod . produce ( uuid4 ( ) . bytes ) <EOL> def test_null_payloads ( self ) : <EOL> """<STR_LIT>""" <EOL> consumer = self . _get_consumer ( ) <EOL> prod = self . _get_producer ( sync = True , min_queued_messages = <NUM_LIT:1> ) <EOL> prod . produce ( None ) <EOL> self . assertIsNone ( consumer . consume ( ) . value ) <EOL> prod . produce ( None , partition_key = b"<STR_LIT>" ) <EOL> self . assertIsNone ( consumer . consume ( ) . value ) <EOL> prod . produce ( b"<STR_LIT>" ) <EOL> self . assertEqual ( b"<STR_LIT>" , consumer . consume ( ) . value ) <EOL> def test_owned_broker_flush_message_larger_then_max_request_size ( self ) : <EOL> """<STR_LIT>""" <EOL> large_payload = b'<STR_LIT>' . join ( [ uuid4 ( ) . bytes for i in range ( <NUM_LIT> ) ] ) <EOL> producer = self . _get_producer ( <EOL> auto_start = False <EOL> ) <EOL> partition = producer . _topic . partitions [ <NUM_LIT:0> ] <EOL> owned_broker = OwnedBroker ( producer , partition . leader , auto_start = False ) <EOL> delivery_report_queue = producer . _cluster . handler . Queue ( ) <EOL> msg = Message ( <EOL> large_payload , <EOL> partition_id = <NUM_LIT:0> , <EOL> delivery_report_q = delivery_report_queue <EOL> ) <EOL> owned_broker . enqueue ( msg ) <EOL> max_request_size = <NUM_LIT:1000> <EOL> assert max_request_size < len ( msg ) <EOL> owned_broker . flush ( <NUM_LIT:0> , max_request_size ) <EOL> q_msg , exc = delivery_report_queue . get ( ) <EOL> assert q_msg is msg <EOL> assert isinstance ( exc , MessageSizeTooLarge ) <EOL> def test_owned_broker_flush_batching_by_max_request_size ( self ) : <EOL> """<STR_LIT>""" <EOL> large_payload = b'<STR_LIT>' . join ( [ uuid4 ( ) . bytes for i in range ( <NUM_LIT> ) ] ) <EOL> producer = self . _get_producer ( <EOL> auto_start = False <EOL> ) <EOL> partition = producer . _topic . partitions [ <NUM_LIT:0> ] <EOL> owned_broker = OwnedBroker ( producer , partition . leader , auto_start = False ) <EOL> for i in range ( <NUM_LIT:100> ) : <EOL> msg = Message ( large_payload , partition_id = <NUM_LIT:0> ) <EOL> owned_broker . enqueue ( msg ) <EOL> batch = owned_broker . flush ( <NUM_LIT:0> , producer . _max_request_size ) <EOL> assert len ( batch ) < <NUM_LIT:100> <EOL> assert sum ( [ len ( m . value ) for m in batch ] ) < producer . _max_request_size <EOL> while batch : <EOL> batch = owned_broker . flush ( <NUM_LIT:0> , producer . _max_request_size ) <EOL> assert len ( batch ) < <NUM_LIT:100> <EOL> assert sum ( [ len ( m . value ) for m in batch ] ) < producer . _max_request_size <EOL> def test_async_produce_compression_large_message ( self ) : <EOL> large_payload = b'<STR_LIT>' . join ( [ uuid4 ( ) . bytes for i in range ( <NUM_LIT:5> ) ] ) <EOL> consumer = self . _get_consumer ( ) <EOL> prod = self . _get_producer ( <EOL> compression = CompressionType . SNAPPY , <EOL> delivery_reports = True <EOL> ) <EOL> prod . produce ( large_payload ) <EOL> report = prod . get_delivery_report ( ) <EOL> self . assertEqual ( report [ <NUM_LIT:0> ] . value , large_payload ) <EOL> self . assertIsNone ( report [ <NUM_LIT:1> ] ) <EOL> message = consumer . consume ( ) <EOL> assert message . value == large_payload <EOL> for i in range ( <NUM_LIT:10> ) : <EOL> prod . produce ( large_payload ) <EOL> reports = [ ] <EOL> def ensure_all_messages_produced ( ) : <EOL> report = prod . get_delivery_report ( ) <EOL> reports . append ( report ) <EOL> assert len ( reports ) == <NUM_LIT:10> <EOL> retry ( ensure_all_messages_produced , retry_time = <NUM_LIT:30> , wait_between_tries = <NUM_LIT:0.5> ) <EOL> for report in reports : <EOL> self . assertEqual ( report [ <NUM_LIT:0> ] . value , large_payload ) <EOL> self . assertIsNone ( report [ <NUM_LIT:1> ] ) <EOL> msgs = [ ] <EOL> def ensure_all_messages_consumed ( ) : <EOL> msg = consumer . consume ( ) <EOL> if msg : <EOL> msgs . append ( msg ) <EOL> assert len ( msgs ) == <NUM_LIT:10> <EOL> retry ( ensure_all_messages_consumed , retry_time = <NUM_LIT:15> ) <EOL> def test_async_produce_large_message ( self ) : <EOL> consumer = self . _get_consumer ( ) <EOL> large_payload = b'<STR_LIT>' . join ( [ uuid4 ( ) . bytes for i in range ( <NUM_LIT> ) ] ) <EOL> assert len ( large_payload ) / <NUM_LIT> / <NUM_LIT> < <NUM_LIT:1.0> <EOL> prod = self . _get_producer ( delivery_reports = True ) <EOL> prod . produce ( large_payload ) <EOL> report = prod . get_delivery_report ( ) <EOL> self . assertEqual ( report [ <NUM_LIT:0> ] . value , large_payload ) <EOL> self . assertIsNone ( report [ <NUM_LIT:1> ] ) <EOL> message = consumer . consume ( ) <EOL> assert message . value == large_payload <EOL> for i in range ( <NUM_LIT:10> ) : <EOL> prod . produce ( large_payload ) <EOL> reports = [ ] <EOL> def ensure_all_messages_produced ( ) : <EOL> report = prod . get_delivery_report ( ) <EOL> reports . append ( report ) <EOL> assert len ( reports ) == <NUM_LIT:10> <EOL> retry ( ensure_all_messages_produced , retry_time = <NUM_LIT:30> , wait_between_tries = <NUM_LIT:0.5> ) <EOL> for report in reports : <EOL> self . assertEqual ( report [ <NUM_LIT:0> ] . value , large_payload ) <EOL> self . assertIsNone ( report [ <NUM_LIT:1> ] ) <EOL> msgs = [ ] <EOL> def ensure_all_messages_consumed ( ) : <EOL> msg = consumer . consume ( ) <EOL> if msg : <EOL> msgs . append ( msg ) <EOL> assert len ( msgs ) == <NUM_LIT:10> <EOL> retry ( ensure_all_messages_consumed , retry_time = <NUM_LIT:15> ) <EOL> @ pytest . mark . skipif ( platform . python_implementation ( ) == "<STR_LIT>" , <EOL> reason = "<STR_LIT>" ) <EOL> class TestGEventProducer ( ProducerIntegrationTests ) : <EOL> USE_GEVENT = True <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> unittest2 . main ( ) </s>
