<s> import socket <EOL> from time import sleep <EOL> from mock import ANY , MagicMock , patch <EOL> import six <EOL> from . import unittest <EOL> from kafka import SimpleClient <EOL> from kafka . conn import KafkaConnection <EOL> from kafka . errors import ( <EOL> KafkaUnavailableError , LeaderNotAvailableError , KafkaTimeoutError , <EOL> UnknownTopicOrPartitionError , ConnectionError , FailedPayloadsError ) <EOL> from kafka . future import Future <EOL> from kafka . protocol import KafkaProtocol , create_message <EOL> from kafka . protocol . metadata import MetadataResponse <EOL> from kafka . structs import ProduceRequestPayload , BrokerMetadata , TopicPartition <EOL> from test . testutil import Timer <EOL> NO_ERROR = <NUM_LIT:0> <EOL> UNKNOWN_TOPIC_OR_PARTITION = <NUM_LIT:3> <EOL> NO_LEADER = <NUM_LIT:5> <EOL> def mock_conn ( conn , success = True ) : <EOL> mocked = MagicMock ( ) <EOL> mocked . connected . return_value = True <EOL> if success : <EOL> mocked . send . return_value = Future ( ) . success ( True ) <EOL> else : <EOL> mocked . send . return_value = Future ( ) . failure ( Exception ( ) ) <EOL> conn . return_value = mocked <EOL> class TestSimpleClient ( unittest . TestCase ) : <EOL> def test_init_with_list ( self ) : <EOL> with patch . object ( SimpleClient , '<STR_LIT>' ) : <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> self . assertEqual ( <EOL> sorted ( [ ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) , ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) , <EOL> ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) ] ) , <EOL> sorted ( client . hosts ) ) <EOL> def test_init_with_csv ( self ) : <EOL> with patch . object ( SimpleClient , '<STR_LIT>' ) : <EOL> client = SimpleClient ( hosts = '<STR_LIT>' ) <EOL> self . assertEqual ( <EOL> sorted ( [ ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) , ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) , <EOL> ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) ] ) , <EOL> sorted ( client . hosts ) ) <EOL> def test_init_with_unicode_csv ( self ) : <EOL> with patch . object ( SimpleClient , '<STR_LIT>' ) : <EOL> client = SimpleClient ( hosts = u'<STR_LIT>' ) <EOL> self . assertEqual ( <EOL> sorted ( [ ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) , ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) , <EOL> ( '<STR_LIT>' , <NUM_LIT> , socket . AF_UNSPEC ) ] ) , <EOL> sorted ( client . hosts ) ) <EOL> @ patch . object ( SimpleClient , '<STR_LIT>' ) <EOL> @ patch . object ( SimpleClient , '<STR_LIT>' ) <EOL> def test_send_broker_unaware_request_fail ( self , load_metadata , conn ) : <EOL> mocked_conns = { <EOL> ( '<STR_LIT>' , <NUM_LIT> ) : MagicMock ( ) , <EOL> ( '<STR_LIT>' , <NUM_LIT> ) : MagicMock ( ) <EOL> } <EOL> for val in mocked_conns . values ( ) : <EOL> mock_conn ( val , success = False ) <EOL> def mock_get_conn ( host , port , afi ) : <EOL> return mocked_conns [ ( host , port ) ] <EOL> conn . side_effect = mock_get_conn <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> req = KafkaProtocol . encode_metadata_request ( ) <EOL> with self . assertRaises ( KafkaUnavailableError ) : <EOL> client . _send_broker_unaware_request ( payloads = [ '<STR_LIT>' ] , <EOL> encoder_fn = MagicMock ( return_value = '<STR_LIT>' ) , <EOL> decoder_fn = lambda x : x ) <EOL> for key , conn in six . iteritems ( mocked_conns ) : <EOL> conn . send . assert_called_with ( '<STR_LIT>' ) <EOL> def test_send_broker_unaware_request ( self ) : <EOL> mocked_conns = { <EOL> ( '<STR_LIT>' , <NUM_LIT> ) : MagicMock ( ) , <EOL> ( '<STR_LIT>' , <NUM_LIT> ) : MagicMock ( ) , <EOL> ( '<STR_LIT>' , <NUM_LIT> ) : MagicMock ( ) <EOL> } <EOL> mock_conn ( mocked_conns [ ( '<STR_LIT>' , <NUM_LIT> ) ] , success = False ) <EOL> mock_conn ( mocked_conns [ ( '<STR_LIT>' , <NUM_LIT> ) ] , success = False ) <EOL> future = Future ( ) <EOL> mocked_conns [ ( '<STR_LIT>' , <NUM_LIT> ) ] . send . return_value = future <EOL> mocked_conns [ ( '<STR_LIT>' , <NUM_LIT> ) ] . recv . side_effect = lambda : future . success ( '<STR_LIT>' ) <EOL> def mock_get_conn ( host , port , afi ) : <EOL> return mocked_conns [ ( host , port ) ] <EOL> with patch . object ( SimpleClient , '<STR_LIT>' ) : <EOL> with patch . object ( SimpleClient , '<STR_LIT>' , side_effect = mock_get_conn ) : <EOL> client = SimpleClient ( hosts = '<STR_LIT>' ) <EOL> resp = client . _send_broker_unaware_request ( payloads = [ '<STR_LIT>' ] , <EOL> encoder_fn = MagicMock ( ) , <EOL> decoder_fn = lambda x : x ) <EOL> self . assertEqual ( '<STR_LIT>' , resp ) <EOL> mocked_conns [ ( '<STR_LIT>' , <NUM_LIT> ) ] . recv . assert_called_once_with ( ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_load_metadata ( self , protocol , conn ) : <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_ERROR , <NUM_LIT:0> , <NUM_LIT:1> , [ <NUM_LIT:1> , <NUM_LIT:2> ] , [ <NUM_LIT:1> , <NUM_LIT:2> ] ) <EOL> ] ) , <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_LEADER , <NUM_LIT:0> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ( NO_LEADER , <NUM_LIT:1> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ] ) , <EOL> ( NO_LEADER , '<STR_LIT>' , [ ] ) , <EOL> ( UNKNOWN_TOPIC_OR_PARTITION , '<STR_LIT>' , [ ] ) , <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_ERROR , <NUM_LIT:0> , <NUM_LIT:0> , [ <NUM_LIT:0> , <NUM_LIT:1> ] , [ <NUM_LIT:0> , <NUM_LIT:1> ] ) , <EOL> ( NO_ERROR , <NUM_LIT:1> , <NUM_LIT:1> , [ <NUM_LIT:1> , <NUM_LIT:0> ] , [ <NUM_LIT:1> , <NUM_LIT:0> ] ) , <EOL> ( NO_ERROR , <NUM_LIT:2> , <NUM_LIT:0> , [ <NUM_LIT:0> , <NUM_LIT:1> ] , [ <NUM_LIT:0> , <NUM_LIT:1> ] ) <EOL> ] ) <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> self . assertDictEqual ( { <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:0> ) : brokers [ <NUM_LIT:1> ] , <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:0> ) : None , <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:1> ) : None , <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:0> ) : brokers [ <NUM_LIT:0> ] , <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:1> ) : brokers [ <NUM_LIT:1> ] , <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:2> ) : brokers [ <NUM_LIT:0> ] } , <EOL> client . topics_to_brokers ) <EOL> with self . assertRaises ( LeaderNotAvailableError ) : <EOL> client . load_metadata_for_topics ( '<STR_LIT>' ) <EOL> with self . assertRaises ( UnknownTopicOrPartitionError ) : <EOL> client . load_metadata_for_topics ( '<STR_LIT>' ) <EOL> client . load_metadata_for_topics ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_has_metadata_for_topic ( self , protocol , conn ) : <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( NO_LEADER , '<STR_LIT>' , [ ] ) , <EOL> ( UNKNOWN_TOPIC_OR_PARTITION , '<STR_LIT>' , [ ] ) , <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_LEADER , <NUM_LIT:0> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ( NO_LEADER , <NUM_LIT:1> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ] ) , <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> self . assertFalse ( client . has_metadata_for_topic ( '<STR_LIT>' ) ) <EOL> self . assertFalse ( client . has_metadata_for_topic ( '<STR_LIT>' ) ) <EOL> self . assertTrue ( client . has_metadata_for_topic ( '<STR_LIT>' ) ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_ensure_topic_exists ( self , decode_metadata_response , conn ) : <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( NO_LEADER , '<STR_LIT>' , [ ] ) , <EOL> ( UNKNOWN_TOPIC_OR_PARTITION , '<STR_LIT>' , [ ] ) , <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_LEADER , <NUM_LIT:0> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ( NO_LEADER , <NUM_LIT:1> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ] ) , <EOL> ] <EOL> decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> with self . assertRaises ( UnknownTopicOrPartitionError ) : <EOL> client . ensure_topic_exists ( '<STR_LIT>' , timeout = <NUM_LIT:1> ) <EOL> with self . assertRaises ( KafkaTimeoutError ) : <EOL> client . ensure_topic_exists ( '<STR_LIT>' , timeout = <NUM_LIT:1> ) <EOL> client . ensure_topic_exists ( '<STR_LIT>' , timeout = <NUM_LIT:1> ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_get_leader_for_partitions_reloads_metadata ( self , protocol , conn ) : <EOL> "<STR_LIT>" <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( NO_LEADER , '<STR_LIT>' , [ ] ) <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> self . assertDictEqual ( { } , client . topics_to_brokers ) <EOL> topics = [ <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_ERROR , <NUM_LIT:0> , <NUM_LIT:0> , [ <NUM_LIT:0> , <NUM_LIT:1> ] , [ <NUM_LIT:0> , <NUM_LIT:1> ] ) <EOL> ] ) <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> leader = client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:0> ) <EOL> self . assertEqual ( brokers [ <NUM_LIT:0> ] , leader ) <EOL> self . assertDictEqual ( { <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:0> ) : brokers [ <NUM_LIT:0> ] } , <EOL> client . topics_to_brokers ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_get_leader_for_unassigned_partitions ( self , protocol , conn ) : <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( NO_LEADER , '<STR_LIT>' , [ ] ) , <EOL> ( UNKNOWN_TOPIC_OR_PARTITION , '<STR_LIT>' , [ ] ) , <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> self . assertDictEqual ( { } , client . topics_to_brokers ) <EOL> with self . assertRaises ( LeaderNotAvailableError ) : <EOL> client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:0> ) <EOL> with self . assertRaises ( UnknownTopicOrPartitionError ) : <EOL> client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:0> ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_get_leader_exceptions_when_noleader ( self , protocol , conn ) : <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_LEADER , <NUM_LIT:0> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ( NO_LEADER , <NUM_LIT:1> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ] ) , <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> self . assertDictEqual ( <EOL> { <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:0> ) : None , <EOL> TopicPartition ( '<STR_LIT>' , <NUM_LIT:1> ) : None <EOL> } , <EOL> client . topics_to_brokers ) <EOL> with self . assertRaises ( LeaderNotAvailableError ) : <EOL> self . assertIsNone ( client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> with self . assertRaises ( LeaderNotAvailableError ) : <EOL> self . assertIsNone ( client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:1> ) ) <EOL> with self . assertRaises ( UnknownTopicOrPartitionError ) : <EOL> self . assertIsNone ( client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:2> ) ) <EOL> topics = [ <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_ERROR , <NUM_LIT:0> , <NUM_LIT:0> , [ <NUM_LIT:0> , <NUM_LIT:1> ] , [ <NUM_LIT:0> , <NUM_LIT:1> ] ) , <EOL> ( NO_ERROR , <NUM_LIT:1> , <NUM_LIT:1> , [ <NUM_LIT:1> , <NUM_LIT:0> ] , [ <NUM_LIT:1> , <NUM_LIT:0> ] ) <EOL> ] ) , <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> self . assertEqual ( brokers [ <NUM_LIT:0> ] , client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> self . assertEqual ( brokers [ <NUM_LIT:1> ] , client . _get_leader_for_partition ( '<STR_LIT>' , <NUM_LIT:1> ) ) <EOL> @ patch . object ( SimpleClient , '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_send_produce_request_raises_when_noleader ( self , protocol , conn ) : <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( NO_ERROR , '<STR_LIT>' , [ <EOL> ( NO_LEADER , <NUM_LIT:0> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ( NO_LEADER , <NUM_LIT:1> , - <NUM_LIT:1> , [ ] , [ ] ) , <EOL> ] ) , <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> requests = [ ProduceRequestPayload ( <EOL> "<STR_LIT>" , <NUM_LIT:0> , <EOL> [ create_message ( "<STR_LIT:a>" ) , create_message ( "<STR_LIT:b>" ) ] ) ] <EOL> with self . assertRaises ( FailedPayloadsError ) : <EOL> client . send_produce_request ( requests ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_send_produce_request_raises_when_topic_unknown ( self , protocol , conn ) : <EOL> mock_conn ( conn ) <EOL> brokers = [ <EOL> BrokerMetadata ( <NUM_LIT:0> , '<STR_LIT>' , <NUM_LIT> ) , <EOL> BrokerMetadata ( <NUM_LIT:1> , '<STR_LIT>' , <NUM_LIT> ) <EOL> ] <EOL> topics = [ <EOL> ( UNKNOWN_TOPIC_OR_PARTITION , '<STR_LIT>' , [ ] ) , <EOL> ] <EOL> protocol . decode_metadata_response . return_value = MetadataResponse [ <NUM_LIT:0> ] ( brokers , topics ) <EOL> client = SimpleClient ( hosts = [ '<STR_LIT>' ] ) <EOL> requests = [ ProduceRequestPayload ( <EOL> "<STR_LIT>" , <NUM_LIT:0> , <EOL> [ create_message ( "<STR_LIT:a>" ) , create_message ( "<STR_LIT:b>" ) ] ) ] <EOL> with self . assertRaises ( FailedPayloadsError ) : <EOL> client . send_produce_request ( requests ) <EOL> def test_timeout ( self ) : <EOL> def _timeout ( * args , ** kwargs ) : <EOL> timeout = args [ <NUM_LIT:1> ] <EOL> sleep ( timeout ) <EOL> raise socket . timeout <EOL> with patch . object ( socket , "<STR_LIT>" , side_effect = _timeout ) : <EOL> with Timer ( ) as t : <EOL> with self . assertRaises ( ConnectionError ) : <EOL> KafkaConnection ( "<STR_LIT>" , <NUM_LIT> , <NUM_LIT:1.0> ) <EOL> self . assertGreaterEqual ( t . interval , <NUM_LIT:1.0> ) <EOL> def test_correlation_rollover ( self ) : <EOL> with patch . object ( SimpleClient , '<STR_LIT>' ) : <EOL> big_num = <NUM_LIT:2> ** <NUM_LIT> - <NUM_LIT:3> <EOL> client = SimpleClient ( hosts = [ ] , correlation_id = big_num ) <EOL> self . assertEqual ( big_num + <NUM_LIT:1> , client . _next_id ( ) ) <EOL> self . assertEqual ( big_num + <NUM_LIT:2> , client . _next_id ( ) ) <EOL> self . assertEqual ( <NUM_LIT:0> , client . _next_id ( ) ) </s>
