From 0f486b7ef3e141296189b1bbc18baf27f046e160 Mon Sep 17 00:00:00 2001
From: Dave Shawley <daves@aweber.com>
Date: Thu, 21 Jan 2016 14:49:11 -0500
Subject: [PATCH] Add s.m.m.testing.FakeInfluxHandler.

---
 docs/api.rst                        |  3 +
 docs/history.rst                    |  1 +
 sprockets/mixins/metrics/testing.py | 89 +++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+)

diff --git a/docs/api.rst b/docs/api.rst
index ff4574e..4761f92 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -56,3 +56,6 @@ contains some helper that make testing a little easier.
 
 .. autoclass:: sprockets.mixins.metrics.testing.FakeStatsdServer
    :members:
+
+.. autoclass:: sprockets.mixins.metrics.testing.FakeInfluxHandler
+   :members:
diff --git a/docs/history.rst b/docs/history.rst
index b11555e..8240969 100644
--- a/docs/history.rst
+++ b/docs/history.rst
@@ -7,5 +7,6 @@ Release History
 ---------------
 - Add :class:`sprockets.mixins.metrics.StatsdMixin`
 - Add :class:`sprockets.mixins.metrics.testing.FakeStatsdServer`
+- Add :class:`sprockets.mixins.metrics.testing.FakeInfluxHandler`
 
 .. _Next Release: https://github.com/sprockets/sprockets.mixins.metrics/compare/0.0.0...master
diff --git a/sprockets/mixins/metrics/testing.py b/sprockets/mixins/metrics/testing.py
index 46e0284..c1b0f39 100644
--- a/sprockets/mixins/metrics/testing.py
+++ b/sprockets/mixins/metrics/testing.py
@@ -1,7 +1,12 @@
+import collections
+import logging
 import re
 import socket
 
+from tornado import gen, web
 
+
+LOGGER = logging.getLogger(__name__)
 STATS_PATTERN = re.compile(r'(?P<path>[^:]*):(?P<value>[^|]*)\|(?P<type>.*)$')
 
 
@@ -81,3 +86,87 @@ class FakeStatsdServer(object):
             raise AssertionError(
                 'Expected metric starting with "{}" in {!r}'.format(
                     prefix, self.datagrams))
+
+
+class FakeInfluxHandler(web.RequestHandler):
+    """
+    Request handler that mimics the InfluxDB write endpoint.
+
+    Install this handler into your testing application and configure
+    the metrics plugin to write to it.  After running a test, you can
+    examine the received measurements by iterating over the
+    ``influx_db`` list in the application object.
+
+    .. code-block:: python
+
+       class TestThatMyStuffWorks(testing.AsyncHTTPTestCase):
+
+           def get_app(self):
+               self.app = web.Application([
+                   web.url('/', HandlerUnderTest),
+                   web.url('/write', metrics.testing.FakeInfluxHandler),
+               ])
+               return self.app
+
+           def setUp(self):
+               super(TestThatMyStuffWorks, self).setUp()
+               self.app.settings[metrics.InfluxDBMixin.SETTINGS_KEY] = {
+                   'measurement': 'stuff',
+                   'write_url': self.get_url('/write'),
+                   'database': 'requests',
+               }
+
+           def test_that_measurements_are_emitted(self):
+               self.fetch('/')  # invokes handler under test
+               measurements = metrics.testing.FakeInfluxHandler(
+                   self.app, 'requests', self)
+               for key, fields, timestamp in measurements:
+                   # inspect measurements
+
+    """
+
+    def initialize(self):
+        super(FakeInfluxHandler, self).initialize()
+        self.logger = LOGGER.getChild(__name__)
+        if not hasattr(self.application, 'influx_db'):
+            self.application.influx_db = collections.defaultdict(list)
+
+    def post(self):
+        db = self.get_query_argument('db')
+        payload = self.request.body.decode('utf-8')
+        for line in payload.splitlines():
+            self.logger.debug('received "%s"', line)
+            key, fields, timestamp = line.split()
+            self.application.influx_db[db].append((key, fields, timestamp))
+        self.set_status(204)
+
+    @staticmethod
+    def get_messages(application, database, test_case):
+        """
+        Wait for measurements to show up and return them.
+
+        :param tornado.web.Application application: application that
+            :class:`.FakeInfluxHandler` is writing to
+        :param str database: database to retrieve
+        :param tornado.testing.AsyncTestCase test_case: test case
+            that is being executed
+        :return: measurements received as a :class:`list` of
+            (key, fields, timestamp) tuples
+
+        Since measurements are sent asynchronously from within the
+        ``on_finish`` handler they are usually not sent by the time
+        that the test case has stopped the IOloop.  This method accounts
+        for this by running the IOloop until measurements have been
+        received.  It will raise an assertion error if measurements
+        are not received in a reasonable number of runs.
+
+        """
+        for _ in range(0, 10):
+            if hasattr(application, 'influx_db'):
+                if application.influx_db[database]:
+                    return application.influx_db[database]
+            test_case.io_loop.add_future(gen.sleep(0.1),
+                                         lambda _: test_case.stop())
+            test_case.wait()
+        else:
+            test_case.fail('Message not published to InfluxDB before timeout')