<s> import sys <EOL> from mock import MagicMock , patch <EOL> from . import unittest <EOL> from kafka import SimpleConsumer , KafkaConsumer , MultiProcessConsumer <EOL> from kafka . errors import ( <EOL> FailedPayloadsError , KafkaConfigurationError , NotLeaderForPartitionError , <EOL> UnknownTopicOrPartitionError ) <EOL> from kafka . structs import ( <EOL> FetchResponsePayload , OffsetAndMessage , OffsetFetchResponsePayload ) <EOL> class TestKafkaConsumer ( unittest . TestCase ) : <EOL> def test_non_integer_partitions ( self ) : <EOL> with self . assertRaises ( AssertionError ) : <EOL> SimpleConsumer ( MagicMock ( ) , '<STR_LIT>' , '<STR_LIT>' , partitions = [ '<STR_LIT:0>' ] ) <EOL> class TestMultiProcessConsumer ( unittest . TestCase ) : <EOL> @ unittest . skipIf ( sys . platform . startswith ( '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> def test_partition_list ( self ) : <EOL> client = MagicMock ( ) <EOL> partitions = ( <NUM_LIT:0> , ) <EOL> with patch . object ( MultiProcessConsumer , '<STR_LIT>' ) as fetch_last_known_offsets : <EOL> MultiProcessConsumer ( client , '<STR_LIT>' , '<STR_LIT>' , partitions = partitions ) <EOL> self . assertEqual ( fetch_last_known_offsets . call_args [ <NUM_LIT:0> ] , ( partitions , ) ) <EOL> self . assertEqual ( client . get_partition_ids_for_topic . call_count , <NUM_LIT:0> ) <EOL> class TestSimpleConsumer ( unittest . TestCase ) : <EOL> def test_simple_consumer_failed_payloads ( self ) : <EOL> client = MagicMock ( ) <EOL> consumer = SimpleConsumer ( client , group = None , <EOL> topic = '<STR_LIT>' , partitions = [ <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> auto_commit = False ) <EOL> def failed_payloads ( payload ) : <EOL> return FailedPayloadsError ( payload ) <EOL> client . send_fetch_request . side_effect = self . fail_requests_factory ( failed_payloads ) <EOL> consumer . get_messages ( <NUM_LIT:5> ) <EOL> def test_simple_consumer_leader_change ( self ) : <EOL> client = MagicMock ( ) <EOL> consumer = SimpleConsumer ( client , group = None , <EOL> topic = '<STR_LIT>' , partitions = [ <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> auto_commit = False ) <EOL> def not_leader ( request ) : <EOL> return FetchResponsePayload ( request . topic , request . partition , <EOL> NotLeaderForPartitionError . errno , - <NUM_LIT:1> , ( ) ) <EOL> client . send_fetch_request . side_effect = self . fail_requests_factory ( not_leader ) <EOL> consumer . get_messages ( <NUM_LIT:20> ) <EOL> self . assertGreaterEqual ( client . reset_topic_metadata . call_count , <NUM_LIT:1> ) <EOL> self . assertGreaterEqual ( client . load_metadata_for_topics . call_count , <NUM_LIT:1> ) <EOL> def test_simple_consumer_unknown_topic_partition ( self ) : <EOL> client = MagicMock ( ) <EOL> consumer = SimpleConsumer ( client , group = None , <EOL> topic = '<STR_LIT>' , partitions = [ <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> auto_commit = False ) <EOL> def unknown_topic_partition ( request ) : <EOL> return FetchResponsePayload ( request . topic , request . partition , <EOL> UnknownTopicOrPartitionError . errno , - <NUM_LIT:1> , ( ) ) <EOL> client . send_fetch_request . side_effect = self . fail_requests_factory ( unknown_topic_partition ) <EOL> with self . assertRaises ( UnknownTopicOrPartitionError ) : <EOL> consumer . get_messages ( <NUM_LIT:20> ) <EOL> def test_simple_consumer_commit_does_not_raise ( self ) : <EOL> client = MagicMock ( ) <EOL> client . get_partition_ids_for_topic . return_value = [ <NUM_LIT:0> , <NUM_LIT:1> ] <EOL> def mock_offset_fetch_request ( group , payloads , ** kwargs ) : <EOL> return [ OffsetFetchResponsePayload ( p . topic , p . partition , <NUM_LIT:0> , b'<STR_LIT>' , <NUM_LIT:0> ) for p in payloads ] <EOL> client . send_offset_fetch_request . side_effect = mock_offset_fetch_request <EOL> def mock_offset_commit_request ( group , payloads , ** kwargs ) : <EOL> raise FailedPayloadsError ( payloads [ <NUM_LIT:0> ] ) <EOL> client . send_offset_commit_request . side_effect = mock_offset_commit_request <EOL> consumer = SimpleConsumer ( client , group = '<STR_LIT>' , <EOL> topic = '<STR_LIT>' , partitions = [ <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> auto_commit = False ) <EOL> consumer . count_since_commit = <NUM_LIT:10> <EOL> self . assertFalse ( consumer . commit ( partitions = [ <NUM_LIT:0> , <NUM_LIT:1> ] ) ) <EOL> def test_simple_consumer_reset_partition_offset ( self ) : <EOL> client = MagicMock ( ) <EOL> def mock_offset_request ( payloads , ** kwargs ) : <EOL> raise FailedPayloadsError ( payloads [ <NUM_LIT:0> ] ) <EOL> client . send_offset_request . side_effect = mock_offset_request <EOL> consumer = SimpleConsumer ( client , group = '<STR_LIT>' , <EOL> topic = '<STR_LIT>' , partitions = [ <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> auto_commit = False ) <EOL> self . assertEqual ( consumer . reset_partition_offset ( <NUM_LIT:0> ) , None ) <EOL> @ staticmethod <EOL> def fail_requests_factory ( error_factory ) : <EOL> def fail_requests ( payloads , ** kwargs ) : <EOL> responses = [ <EOL> FetchResponsePayload ( payloads [ <NUM_LIT:0> ] . topic , payloads [ <NUM_LIT:0> ] . partition , <NUM_LIT:0> , <NUM_LIT:0> , <EOL> [ OffsetAndMessage ( <EOL> payloads [ <NUM_LIT:0> ] . offset + i , <EOL> "<STR_LIT>" % ( payloads [ <NUM_LIT:0> ] . offset + i ) ) <EOL> for i in range ( <NUM_LIT:10> ) ] ) , <EOL> ] <EOL> for failure in payloads [ <NUM_LIT:1> : ] : <EOL> responses . append ( error_factory ( failure ) ) <EOL> return responses <EOL> return fail_requests </s>
