Separate default handler code for re-use and add test case for custom handler (refs prometheus#120)

rossigee · rossigee · commit 814e24d562ed · 2017-01-25T10:35:47.000+07:00
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -10,6 +10,9 @@
 from wsgiref.simple_server import make_server
 
 from . import core
+
+from handlers.base import handler as default_handler
+
 try:
     from BaseHTTPServer import BaseHTTPRequestHandler
     from BaseHTTPServer import HTTPServer
@@ -222,17 +225,11 @@ def _use_gateway(method, gateway, job, registry, grouping_key, timeout, handler)
     url = url + ''.join(['/{0}/{1}'.format(quote_plus(str(k)), quote_plus(str(v)))
                              for k, v in sorted(grouping_key.items())])
 
-    request = Request(url, data=data)
-    request.add_header('Content-Type', CONTENT_TYPE_LATEST)
-    request.get_method = lambda: method
+    headers=[('Content-Type', CONTENT_TYPE_LATEST)]
     if handler is None:
-        resp = build_opener(handler).open(request, timeout=timeout)
-        if resp.code >= 400:
-            raise IOError("error talking to pushgateway: {0} {1}".format(
-                resp.code, resp.msg))
-    else:
-        handler(url=url, method=lambda: method, timeout=timeout,
-                headers=[('Content-Type', CONTENT_TYPE_LATEST)], content=data)
+        handler = default_handler
+    handler(url=url, method=method, timeout=timeout,
+            headers=headers, data=data)
 
 def instance_ip_grouping_key():
     '''Grouping key with instance set to the IP Address of this host.'''
diff --git a/prometheus_client/handlers/__init__.py b/prometheus_client/handlers/__init__.py
diff --git a/prometheus_client/handlers/base.py b/prometheus_client/handlers/base.py
@@ -0,0 +1,18 @@
+#!/usr/bin/python
+
+try:
+    from urllib2 import build_opener, Request, HTTPHandler
+except ImportError:
+    # Python 3
+    from urllib.request import build_opener, Request, HTTPHandler
+
+def handler(url, method, timeout, headers, data):
+    '''Default handler that implements HTTP/HTTPS connections.'''
+    request = Request(url, data=data)
+    request.get_method = lambda: method
+    for k, v in headers:
+        request.add_header(k, v)
+    resp = build_opener(HTTPHandler).open(request, timeout=timeout)
+    if resp.code >= 400:
+        raise IOError("error talking to pushgateway: {0} {1}".format(
+            resp.code, resp.msg))
diff --git a/tests/test_exposition.py b/tests/test_exposition.py
@@ -13,6 +13,7 @@
 from prometheus_client import CollectorRegistry, generate_latest
 from prometheus_client import push_to_gateway, pushadd_to_gateway, delete_from_gateway
 from prometheus_client import CONTENT_TYPE_LATEST, instance_ip_grouping_key
+from prometheus_client.handlers.base import handler as default_handler
 
 try:
     from BaseHTTPServer import BaseHTTPRequestHandler
@@ -22,7 +23,6 @@
     from http.server import BaseHTTPRequestHandler
     from http.server import HTTPServer
 
-
 class TestGenerateText(unittest.TestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
@@ -165,6 +165,17 @@ def test_delete_with_groupingkey(self):
         self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
         self.assertEqual(self.requests[0][1], b'')
 
+    def test_push_with_handler(self):
+        def my_test_handler(url, method, timeout, headers, data):
+            headers.append(['X-Test-Header', 'foobar'])
+            default_handler(url, method, timeout, headers, data)
+        push_to_gateway(self.address, "my_job", self.registry, handler=my_test_handler)
+        self.assertEqual(self.requests[0][0].command, 'PUT')
+        self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job')
+        self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
+        self.assertEqual(self.requests[0][0].headers.get('x-test-header'), 'foobar')
+        self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
+
     @unittest.skipIf(
         sys.platform == "darwin",
         "instance_ip_grouping_key() does not work on macOS."
