<s> import urllib2 <EOL> try : <EOL> from xml . etree import ElementTree <EOL> except ImportError : <EOL> ElementTree = None <EOL> from test import CollectorTestCase <EOL> from test import get_collector_config <EOL> from test import run_only <EOL> from test import unittest <EOL> from mock import patch <EOL> from diamond . collector import Collector <EOL> from kafkastat import KafkaCollector <EOL> def run_only_if_ElementTree_is_available ( func ) : <EOL> try : <EOL> from xml . etree import ElementTree <EOL> except ImportError : <EOL> ElementTree = None <EOL> pred = lambda : ElementTree is not None <EOL> return run_only ( func , pred ) <EOL> class TestKafkaCollector ( CollectorTestCase ) : <EOL> def setUp ( self ) : <EOL> config = get_collector_config ( '<STR_LIT>' , { <EOL> '<STR_LIT>' : <NUM_LIT:10> <EOL> } ) <EOL> self . collector = KafkaCollector ( config , None ) <EOL> def _get_xml_fixture ( self , name ) : <EOL> fixture = self . getFixture ( name ) <EOL> return ElementTree . fromstring ( fixture . getvalue ( ) ) <EOL> def test_import ( self ) : <EOL> self . assertTrue ( KafkaCollector ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_get ( self , urlopen_mock ) : <EOL> urlopen_mock . return_value = self . getFixture ( '<STR_LIT>' ) <EOL> result = self . collector . _get ( '<STR_LIT>' ) <EOL> result_string = ElementTree . tostring ( result ) <EOL> self . assertEqual ( result_string , '<STR_LIT>' ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_get_httperror ( self , urlopen_mock ) : <EOL> urlopen_mock . side_effect = urllib2 . URLError ( '<STR_LIT>' ) <EOL> result = self . collector . _get ( '<STR_LIT>' ) <EOL> self . assertFalse ( result ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch ( '<STR_LIT>' ) <EOL> def test_get_bad_xml ( self , urlopen_mock ) : <EOL> urlopen_mock . return_value = self . getFixture ( '<STR_LIT>' ) <EOL> result = self . collector . _get ( '<STR_LIT>' ) <EOL> self . assertFalse ( result ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch . object ( KafkaCollector , '<STR_LIT>' ) <EOL> def test_get_mbeans ( self , get_mock ) : <EOL> get_mock . return_value = self . _get_xml_fixture ( '<STR_LIT>' ) <EOL> expected_names = set ( [ <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> ] ) <EOL> found_beans = self . collector . get_mbeans ( '<STR_LIT:*>' ) <EOL> self . assertEqual ( found_beans , expected_names ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch . object ( KafkaCollector , '<STR_LIT>' ) <EOL> def test_get_mbeans_get_fail ( self , get_mock ) : <EOL> get_mock . return_value = None <EOL> found_beans = self . collector . get_mbeans ( '<STR_LIT:*>' ) <EOL> self . assertEqual ( found_beans , None ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch . object ( KafkaCollector , '<STR_LIT>' ) <EOL> def test_query_mbean ( self , get_mock ) : <EOL> get_mock . return_value = self . _get_xml_fixture ( '<STR_LIT>' ) <EOL> expected_metrics = { <EOL> '<STR_LIT>' : long ( '<STR_LIT>' ) , <EOL> '<STR_LIT>' : long ( '<STR_LIT>' ) , <EOL> '<STR_LIT>' : int ( '<STR_LIT>' ) , <EOL> '<STR_LIT>' : long ( '<STR_LIT>' ) , <EOL> } <EOL> metrics = self . collector . query_mbean ( '<STR_LIT>' ) <EOL> self . assertEqual ( metrics , expected_metrics ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch . object ( KafkaCollector , '<STR_LIT>' ) <EOL> def test_query_mbean_with_prefix ( self , get_mock ) : <EOL> get_mock . return_value = self . _get_xml_fixture ( '<STR_LIT>' ) <EOL> expected_metrics = { <EOL> '<STR_LIT>' : long ( '<STR_LIT>' ) , <EOL> '<STR_LIT>' : long ( '<STR_LIT>' ) , <EOL> '<STR_LIT>' : int ( '<STR_LIT>' ) , <EOL> '<STR_LIT>' : long ( '<STR_LIT>' ) , <EOL> } <EOL> metrics = self . collector . query_mbean ( '<STR_LIT>' , <EOL> '<STR_LIT>' ) <EOL> self . assertEqual ( metrics , expected_metrics ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch . object ( KafkaCollector , '<STR_LIT>' ) <EOL> def test_query_mbean_fail ( self , get_mock ) : <EOL> get_mock . return_value = None <EOL> metrics = self . collector . query_mbean ( '<STR_LIT>' ) <EOL> self . assertEqual ( metrics , None ) <EOL> @ run_only_if_ElementTree_is_available <EOL> @ patch ( '<STR_LIT>' ) <EOL> @ patch . object ( Collector , '<STR_LIT>' ) <EOL> def test ( self , publish_mock , urlopen_mock ) : <EOL> urlopen_mock . side_effect = [ <EOL> self . getFixture ( '<STR_LIT>' ) , <EOL> self . getFixture ( '<STR_LIT>' ) , <EOL> self . getFixture ( '<STR_LIT>' ) , <EOL> self . getFixture ( '<STR_LIT>' ) , <EOL> self . getFixture ( '<STR_LIT>' ) , <EOL> self . getFixture ( '<STR_LIT>' ) , <EOL> self . getFixture ( '<STR_LIT>' ) , <EOL> ] <EOL> self . collector . collect ( ) <EOL> expected_metrics = { <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT:0> , <EOL> '<STR_LIT>' : <NUM_LIT:0> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : <NUM_LIT:2> , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> } <EOL> self . assertPublishedMany ( publish_mock , expected_metrics ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> unittest . main ( ) </s>
