8000 Separate default handler code for re-use and add test case for custom… · asv/client_python@c37a15c · GitHub
[go: up one dir, main page]

Skip to content

Commit c37a15c

Browse files
committed
Separate default handler code for re-use and add test case for custom handler (refs prometheus#120)
1 parent 0be1ec5 commit c37a15c

File tree

4 files changed

+37
-11
lines changed

4 files changed

+37
-11
lines changed

prometheus_client/exposition.py

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@
1010
from wsgiref.simple_server import make_server
1111

1212
from . import core
13+
14+
from handlers.base import handler as default_handler
15+
1316
try:
1417
from BaseHTTPServer import BaseHTTPRequestHandler
1518
from BaseHTTPServer import HTTPServer
@@ -222,17 +225,11 @@ def _use_gateway(method, gateway, job, registry, grouping_key, timeout, handler)
222225
url = url + ''.join(['/{0}/{1}'.format(quote_plus(str(k)), quote_plus(str(v)))
223226
for k, v in sorted(grouping_key.items())])
224227

225-
request = Request(url, data=data)
226-
request.add_header('Content-Type', CONTENT_TYPE_LATEST)
227-
request.get_method = lambda: method
228+
headers=[('Content-Type', CONTENT_TYPE_LATEST)]
228229
if handler is None:
229-
resp = build_opener(handler).open(request, timeout=timeout)
230-
if resp.code >= 400:
231-
raise IOError("error talking to pushgateway: {0} {1}".format(
232-
resp.code, resp.msg))
233-
else:
234-
handler(url=url, method=lambda: method, timeout=timeout,
235-
headers=[('Content-Type', CONTENT_TYPE_LATEST)], content=data)
230+
handler = default_handler
231+
handler(url=url, method=method, timeout=timeout,
232+
headers=headers, data=data)
236233

237234
def instance_ip_grouping_key():
238235
'''Grouping key with instance set to the IP Address of this host.'''

prometheus_client/handlers/__init__.py

Whitespace-only changes.

prometheus_client/handlers/base.py

Lines changed: 18 additions & 0 deletions
+
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
#!/usr/bin/python
2
3+
try:
4+
from urllib2 import build_opener, Request, HTTPHandler
5+
except ImportError:
6+
# Python 3
7+
from urllib.request import build_opener, Request, HTTPHandler
8+
9+
def handler(url, method, timeout, headers, data):
10+
'''Default handler that implements HTTP/HTTPS connections.'''
11+
request = Request(url, data=data)
12+
request.get_method = lambda: method
13+
for k, v in headers:
14+
request.add_header(k, v)
15+
resp = build_opener(HTTPHandler).open(request, timeout=timeout)
16+
if resp.code >= 400:
17+
raise IOError("error talking to pushgateway: {0} {1}".format(
18+
resp.code, resp.msg))

tests/test_exposition.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from prometheus_client import CollectorRegistry, generate_latest
1414
from prometheus_client import push_to_gateway, pushadd_to_gateway, delete_from_gateway
1515
from prometheus_client import CONTENT_TYPE_LATEST, instance_ip_grouping_key
16+
from prometheus_client.handlers.base import handler as default_handler
1617

1718
try:
1819
from BaseHTTPServer import BaseHTTPRequestHandler
@@ -22,7 +23,6 @@
2223
from http.server import BaseHTTPRequestHandler
2324
from http.server import HTTPServer
2425

25-
2626
class TestGenerateText(unittest.TestCase):
2727
def setUp(self):
2828
self.registry = CollectorRegistry()
@@ -165,6 +165,17 @@ def test_delete_with_groupingkey(self):
165165
self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
166166
self.assertEqual(self.requests[0][1], b'')
167167

168+
def test_push_with_handler(self):
169+
def my_test_handler(url, method, timeout, headers, data):
170+
headers.append(['X-Test-Header', 'foobar'])
171+
default_handler(url, method, timeout, headers, data)
172+
push_to_gateway(self.address, "my_job", self.registry, handler=my_test_handler)
173+
self.assertEqual(self.requests[0][0].command, 'PUT')
174+
self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job')
175+
self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
176+
self.assertEqual(self.requests[0][0].headers.get('x-test-header'), 'foobar')
177+
self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
178+
168179
@unittest.skipIf(
169180
sys.platform == "darwin",
170181
"instance_ip_grouping_key() does not work on macOS."

0 commit comments

Comments
 (0)
0