<s> from __future__ import absolute_import <EOL> import pytest <EOL> from kafka . coordinator . assignors . range import RangePartitionAssignor <EOL> from kafka . coordinator . assignors . roundrobin import RoundRobinPartitionAssignor <EOL> from kafka . coordinator . protocol import ( <EOL> ConsumerProtocolMemberMetadata , ConsumerProtocolMemberAssignment ) <EOL> @ pytest . fixture <EOL> def cluster ( mocker ) : <EOL> cluster = mocker . MagicMock ( ) <EOL> cluster . partitions_for_topic . return_value = set ( [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> ] ) <EOL> return cluster <EOL> def test_assignor_roundrobin ( cluster ) : <EOL> assignor = RoundRobinPartitionAssignor <EOL> member_metadata = { <EOL> '<STR_LIT>' : assignor . metadata ( set ( [ '<STR_LIT>' , '<STR_LIT>' ] ) ) , <EOL> '<STR_LIT>' : assignor . metadata ( set ( [ '<STR_LIT>' , '<STR_LIT>' ] ) ) , <EOL> } <EOL> ret = assignor . assign ( cluster , member_metadata ) <EOL> expected = { <EOL> '<STR_LIT>' : ConsumerProtocolMemberAssignment ( <EOL> assignor . version , [ ( '<STR_LIT>' , [ <NUM_LIT:0> , <NUM_LIT:2> ] ) , ( '<STR_LIT>' , [ <NUM_LIT:1> ] ) ] , b'<STR_LIT>' ) , <EOL> '<STR_LIT>' : ConsumerProtocolMemberAssignment ( <EOL> assignor . version , [ ( '<STR_LIT>' , [ <NUM_LIT:1> ] ) , ( '<STR_LIT>' , [ <NUM_LIT:0> , <NUM_LIT:2> ] ) ] , b'<STR_LIT>' ) <EOL> } <EOL> assert ret == expected <EOL> assert set ( ret ) == set ( expected ) <EOL> for member in ret : <EOL> assert ret [ member ] . encode ( ) == expected [ member ] . encode ( ) <EOL> def test_assignor_range ( cluster ) : <EOL> assignor = RangePartitionAssignor <EOL> member_metadata = { <EOL> '<STR_LIT>' : assignor . metadata ( set ( [ '<STR_LIT>' , '<STR_LIT>' ] ) ) , <EOL> '<STR_LIT>' : assignor . metadata ( set ( [ '<STR_LIT>' , '<STR_LIT>' ] ) ) , <EOL> } <EOL> ret = assignor . assign ( cluster , member_metadata ) <EOL> expected = { <EOL> '<STR_LIT>' : ConsumerProtocolMemberAssignment ( <EOL> assignor . version , [ ( '<STR_LIT>' , [ <NUM_LIT:0> , <NUM_LIT:1> ] ) , ( '<STR_LIT>' , [ <NUM_LIT:0> , <NUM_LIT:1> ] ) ] , b'<STR_LIT>' ) , <EOL> '<STR_LIT>' : ConsumerProtocolMemberAssignment ( <EOL> assignor . version , [ ( '<STR_LIT>' , [ <NUM_LIT:2> ] ) , ( '<STR_LIT>' , [ <NUM_LIT:2> ] ) ] , b'<STR_LIT>' ) <EOL> } <EOL> assert ret == expected <EOL> assert set ( ret ) == set ( expected ) <EOL> for member in ret : <EOL> assert ret [ member ] . encode ( ) == expected [ member ] . encode ( ) </s>
