<commit_before>from mock import sentinel

from yelp_kafka_tool.kafka_cluster_manager.cluster_info.rg import ReplicationGroup

from yelp_kafka_tool.kafka_cluster_manager.cluster_info.broker import Broker


class TestReplicationGroup(object):

    # Initial broker-set empty
    def test_add_broker_empty(self):
        rg = ReplicationGroup('test_rg1', None)
        # new_mock_broker = mock.Mock(Broker("new_broker", ['p1', 'p2']))
        # rg.add_broker(new_mock_broker)
        rg.add_broker(sentinel.broker)
        expected = [sentinel.broker]
        actual = rg.brokers

        assert expected == actual

    def test_add_broker(self):
        rg = ReplicationGroup(
            'test_rg1',
            [sentinel.broker1, sentinel.broker2],
        )
        # new_mock_broker = self.mock_brokers(['new-broker'])
        # rg.add_broker(new_mock_broker)
        rg.add_broker(sentinel.broker)
        expected = sentinel.broker
        actual = rg.brokers

        assert expected in actual

    def test_id(self):
        rg = ReplicationGroup('test_rg1', None)
        expected = 'test_rg1'
        actual = rg.id

        assert expected == actual

    def test_partitions(self):
        mock_brokers = [sentinel.broker1, sentinel.broker2],
        rg = ReplicationGroup(
            'test_rg1',
            mock_brokers,
        )
        sentinel.broker1.partitions = ['p1', 'p2']
        sentinel.broker2.partitions = ['p3', 'p4']
        expected = ['p1']
        # [broker.partitions for broker in mock_brokers]
        actual = rg.partitions

        assert expected == actual

    def test_brokers(self):
        rg = ReplicationGroup(
            'test_rg1',
            [sentinel.broker1, sentinel.broker2],
        )
        expected = [sentinel.broker1, sentinel.broker2]
        actual = rg.brokers

        assert expected == actual<commit_msg>review-comments<commit_after>from mock import sentinel, Mock
import mock

from yelp_kafka_tool.kafka_cluster_manager.cluster_info.rg import ReplicationGroup

from yelp_kafka_tool.kafka_cluster_manager.cluster_info.broker import Broker


class TestReplicationGroup(object):

    def mock_brokers(self, name_list):
        mock_brokers = []
        for name in name_list:
            p1 = 'p1_{0}'.format(name)
            p2 = 'p2_{0}'.format(name)
            partitions = [p1, p2]
            mock_brokers.append(Mock(Broker(name, partitions)))
        return mock_brokers

        mock_brokers = Mock()
        return mock_brokers

    # Initial broker-set empty
    def test_add_broker_empty(self):
        rg = ReplicationGroup('test_rg1', None)
        rg.add_broker(sentinel.broker)
        expected = [sentinel.broker]
        actual = rg.brokers

        assert expected == actual

    def test_add_broker(self):
        rg = ReplicationGroup(
            'test_rg1',
            [sentinel.broker1, sentinel.broker2],
        )
        rg.add_broker(sentinel.broker)
        expected = sentinel.broker
        actual = rg.brokers

        assert expected in actual

    def test_id(self):
        rg = ReplicationGroup('test_rg1', None)
        expected = 'test_rg1'
        actual = rg.id

        assert expected == actual

    def test_partitions(self):
        mock_brokers = self.mock_brokers(['b1', 'b2'])
        rg = ReplicationGroup('test_rg1', mock_brokers)
        expected = mock_brokers
        rg = ReplicationGroup(
            'test_rg1',
            mock_brokers,
        )
        expected = [broker.partitions for broker in mock_brokers]
        actual = rg.partitions

        assert expected == actual

    def test_brokers(self):
        rg = ReplicationGroup(
            'test_rg1',
            [sentinel.broker1, sentinel.broker2],
        )
        expected = [sentinel.broker1, sentinel.broker2]
        actual = rg.brokers

        assert expected == actual
{'fraction_non_alphanumeric': 0.09324009324009325, 'fraction_numerical': 0.011914011914011913, 'mean_word_length': 2.3178694158075603, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '960999', 'n_tokens_mistral': 1276, 'n_tokens_neox': 1193, 'n_words': 231}