From 8a22c4340feca688ca854070edfac2ce7588a0c2 Mon Sep 17 00:00:00 2001 From: Emil Madsen Date: Mon, 17 Feb 2020 20:37:31 +0100 Subject: [PATCH] Test WSGI, utilize _bake_output in MetricsHandler Signed-off-by: Emil Madsen --- prometheus_client/asgi.py | 8 +++-- prometheus_client/exposition.py | 26 +++++++-------- tests/test_wsgi.py | 58 +++++++++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 16 deletions(-) create mode 100644 tests/test_wsgi.py diff --git a/prometheus_client/asgi.py b/prometheus_client/asgi.py index e1636e6a..5c6505f2 100644 --- a/prometheus_client/asgi.py +++ b/prometheus_client/asgi.py @@ -9,12 +9,15 @@ def make_asgi_app(registry=REGISTRY): async def prometheus_app(scope, receive, send): assert scope.get("type") == "http" + # Prepare parameters params = parse_qs(scope.get('query_string', b'')) accept_header = "Accept: " + ",".join([ value.decode("utf8") for (name, value) in scope.get('headers') if name.decode("utf8") == 'accept' ]) - status, headers, output = _bake_output(registry, accept_header, params) + # Bake output + status, header, output = _bake_output(registry, accept_header, params) + # Return output payload = await receive() if payload.get("type") == "http.request": await send( @@ -22,8 +25,7 @@ async def prometheus_app(scope, receive, send): "type": "http.response.start", "status": int(status.split(' ')[0]), "headers": [ - [key.encode('utf8'), value.encode('utf8')] - for (key,value) in headers + (x.encode('utf8') for x in header) ] } ) diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py index 4e32ff8f..1fb97f63 100644 --- a/prometheus_client/exposition.py +++ b/prometheus_client/exposition.py @@ -38,17 +38,20 @@ def _bake_output(registry, accept_header, params): if 'name[]' in params: registry = registry.restricted_registry(params['name[]']) output = encoder(registry) - return str('200 OK'), [(str('Content-type'), content_type)], output + return str('200 OK'), (str('Content-Type'), content_type), output def make_wsgi_app(registry=REGISTRY): """Create a WSGI app which serves the metrics from a registry.""" def prometheus_app(environ, start_response): + # Prepare parameters accept_header = environ.get('HTTP_ACCEPT') params = parse_qs(environ.get('QUERY_STRING', '')) - status, headers, output = _bake_output(registry, accept_header, params) - start_response(status, headers) + # Bake output + status, header, output = _bake_output(registry, accept_header, params) + # Return output + start_response(status, [header]) return [output] return prometheus_app @@ -147,18 +150,15 @@ class MetricsHandler(BaseHTTPRequestHandler): registry = REGISTRY def do_GET(self): + # Prepare parameters registry = self.registry + accept_header = self.headers.get('Accept') params = parse_qs(urlparse(self.path).query) - encoder, content_type = choose_encoder(self.headers.get('Accept')) - if 'name[]' in params: - registry = registry.restricted_registry(params['name[]']) - try: - output = encoder(registry) - except: - self.send_error(500, 'error generating metric output') - raise - self.send_response(200) - self.send_header('Content-Type', content_type) + # Bake output + status, header, output = _bake_output(registry, accept_header, params) + # Return output + self.send_response(int(status.split(' ')[0])) + self.send_header('Content-Type', header[1]) self.end_headers() self.wfile.write(output) diff --git a/tests/test_wsgi.py b/tests/test_wsgi.py new file mode 100644 index 00000000..8936ed4f --- /dev/null +++ b/tests/test_wsgi.py @@ -0,0 +1,58 @@ +from __future__ import absolute_import, unicode_literals + +import sys + +from prometheus_client import CollectorRegistry, Counter, generate_latest +from prometheus_client.exposition import CONTENT_TYPE_LATEST + +if sys.version_info < (2, 7): + from unittest2 import skipUnless +else: + from unittest import skipUnless + +from prometheus_client import make_wsgi_app +from unittest import TestCase +from wsgiref.util import setup_testing_defaults +from parameterized import parameterized + + +class WSGITest(TestCase): + def setUp(self): + self.registry = CollectorRegistry() + self.captured_status = None + self.captured_headers = None + + def capture(self, status, header): + self.captured_status = status + self.captured_headers = header + + @parameterized.expand([ + ["counter", "A counter"], + ["counter", "Another counter"], + ["requests", "Number of requests"], + ["failed_requests", "Number of failed requests"], + ]) + def test_reports_metrics(self, metric_name, help_text): + """ + WSGI app serves the metrics from the provided registry. + """ + c = Counter(metric_name, help_text, registry=self.registry) + c.inc() + # Setup WSGI environment + environ = {} + setup_testing_defaults(environ) + # Create and run WSGI app + app = make_wsgi_app(self.registry) + outputs = app(environ, self.capture) + # Assert outputs + self.assertEqual(len(outputs), 1) + output = outputs[0].decode('utf8') + # Status code + self.assertEqual(self.captured_status, "200 OK") + # Headers + self.assertEqual(len(self.captured_headers), 1) + self.assertEqual(self.captured_headers[0], ("Content-Type", CONTENT_TYPE_LATEST)) + # Body + self.assertIn("# HELP " + metric_name + "_total " + help_text + "\n", output) + self.assertIn("# TYPE " + metric_name + "_total counter\n", output) + self.assertIn(metric_name + "_total 1.0\n", output)