Test WSGI, utilize _bake_output in MetricsHandler

Signed-off-by: Emil Madsen <sovende@gmail.com>

prometheus_client/asgi.py

@@ -9,21 +9,23 @@ 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(
                 {
                     "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)
                     ]
                 }
             )

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)
 

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)