From 0bc2f8196d98a792449c13c6cfeee18ed688d09b Mon Sep 17 00:00:00 2001 From: Kurt Scherer Date: Thu, 23 Aug 2018 10:07:07 -0500 Subject: [PATCH] Wrap ipv6 addr in brackets for prometheus sizer In the Prometheus sizer, the url is interpolated with the hostname, which can be an ipv6 address and must be wrapped in brackets to separate it from the port number. This commit checks if the hostname contains 2 or more colons and adds brackets if they don't already exist. Fixes #96. --- kafka/tools/assigner/sizers/prometheus.py | 3 +++ tests/tools/assigner/sizers/test_prometheus.py | 17 +++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/kafka/tools/assigner/sizers/prometheus.py b/kafka/tools/assigner/sizers/prometheus.py index 457ab8d..1eb41c3 100644 --- a/kafka/tools/assigner/sizers/prometheus.py +++ b/kafka/tools/assigner/sizers/prometheus.py @@ -83,6 +83,9 @@ def _parse_prometheus_metric(self, text): return metric or None def _get_prometheus_metrics(self, hostname, port, path): + # When hostname is an ipv6 addr, add brackets to separate addr from port + if hostname.count(':') >= 2 and hostname[0] != '[' and hostname[-1] != ']': + hostname = '[{}]'.format(hostname) url = 'http://{}:{}{}'.format(hostname, port, path) body = '' try: diff --git a/tests/tools/assigner/sizers/test_prometheus.py b/tests/tools/assigner/sizers/test_prometheus.py index a2f4432..5ee1e75 100644 --- a/tests/tools/assigner/sizers/test_prometheus.py +++ b/tests/tools/assigner/sizers/test_prometheus.py @@ -159,6 +159,23 @@ def test_sizer_get_prometheus_metrics_200(self, mock_urlopen): metrics = sizer._get_prometheus_metrics('localhost', 1234, '/metrics') self.assertSequenceEqual(metrics, ['404 not found']) + @patch('kafka.tools.assigner.sizers.prometheus.urlopen') + def test_sizer_get_prometheus_metrics_ipv6(self, mock_urlopen): + test_paths = [ + ('1.2.3.4', '1.2.3.4'), + ('username:password@example.com', 'username:password@example.com'), + ('::1', '[::1]'), + ('[2001:db8::1]', '[2001:db8::1]'), + ] + for before, after in test_paths: + m = Mock() + m.getcode.return_value = 200 + mock_urlopen.return_value = m + sizer = SizerPrometheus(self.args, self.cluster) + sizer._get_prometheus_metrics(before, 1234, '/metrics') + mock_urlopen.assert_called_once_with('http://{}:1234/metrics'.format(after)) + mock_urlopen.reset_mock() + @patch('kafka.tools.assigner.sizers.prometheus.urlopen') def test_sizer_get_prometheus_metrics_404(self, mock_urlopen): m = Mock()