From 4c26834f040e094d2b844fbf18da6758bf0bf270 Mon Sep 17 00:00:00 2001
From: Robin Klingsberg <robin@ringingmountain.com>
Date: Sat, 16 Mar 2019 16:07:27 -0400
Subject: [PATCH] Increase test coverage

---
 HISTORY.rst                            |  1 +
 sprockets/mixins/correlation/mixins.py |  2 +-
 tests.py                               | 90 +++++++++++++++++++++++++-
 3 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index 2b0655e..3182a21 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -5,6 +5,7 @@ Version History
 ~~~~~~~~~~~~~~~~~~~~~~
 - Add support for Tornado 6
 - Move requirements files to requires/
+- Increase test coverage
 
 `2.0.0`_ (26-Nov-2018)
 ~~~~~~~~~~~~~~~~~~~~~~
diff --git a/sprockets/mixins/correlation/mixins.py b/sprockets/mixins/correlation/mixins.py
index 7703ec6..fba8d6c 100644
--- a/sprockets/mixins/correlation/mixins.py
+++ b/sprockets/mixins/correlation/mixins.py
@@ -46,7 +46,7 @@ class HandlerMixin(object):
         # one exists.  We also want to set it in the outgoing response
         # which the property setter does for us.
         maybe_future = super(HandlerMixin, self).prepare()
-        if concurrent.is_future(maybe_future):
+        if concurrent.is_future(maybe_future):  # pragma: nocover
             await maybe_future
 
         correlation_id = self.get_request_header(self.__header_name, None)
diff --git a/tests.py b/tests.py
index 72ac07f..a279fff 100644
--- a/tests.py
+++ b/tests.py
@@ -1,20 +1,28 @@
 import uuid
 import unittest
+import unittest.mock
 
 from tornado import testing, web
 
 from sprockets.mixins import correlation
+from sprockets.mixins.correlation.mixins import correlation_id_logger
 
 
-class CorrelatedRequestHandler(correlation.HandlerMixin, web.RequestHandler):
+class RequestHandler(web.RequestHandler):
 
     def get(self, status_code):
         status_code = int(status_code)
+        self.set_status(status_code)
+
         if status_code >= 300:
             raise web.HTTPError(status_code)
         self.write('status {0}'.format(status_code))
 
 
+class CorrelatedRequestHandler(correlation.HandlerMixin, RequestHandler):
+    pass
+
+
 class CorrelationMixinTests(testing.AsyncHTTPTestCase):
 
     def get_app(self):
@@ -35,3 +43,83 @@ class CorrelationMixinTests(testing.AsyncHTTPTestCase):
         response = self.fetch('/status/500',
                               headers={'Correlation-Id': correlation_id})
         self.assertEqual(response.headers['correlation-id'], correlation_id)
+
+
+class CorrelationIDLoggerTests(testing.AsyncHTTPTestCase):
+
+    def get_app(self):
+        return web.Application([
+            (r'/status/(?P<status_code>\d+)', CorrelatedRequestHandler),
+            (r'/status/no-correlation/(?P<status_code>\d+)', RequestHandler),
+        ], log_function=correlation_id_logger)
+
+    def setUp(self):
+        self.patcher = unittest.mock.patch(
+            'sprockets.mixins.correlation.mixins.log.access_log')
+        self.access_logger = self.patcher.start()
+        super(CorrelationIDLoggerTests, self).setUp()
+
+    def tearDown(self):
+        self.patcher.stop()
+        super(CorrelationIDLoggerTests, self).tearDown()
+
+    def test_lt_400_logs_info(self):
+        for status in (200, 202):
+            response = self.fetch('/status/{}'.format(status))
+
+            self.access_logger.info.assert_any_call(
+                "%d %s %.2fms {CID %s}",
+                status,
+                unittest.mock.ANY,
+                unittest.mock.ANY,
+                response.headers['correlation-id']
+            )
+
+    def test_gte_400_lt_500_logs_warning(self):
+        for status in (400, 429):
+            response = self.fetch('/status/{}'.format(status))
+
+            self.access_logger.warning.assert_any_call(
+                "%d %s %.2fms {CID %s}",
+                status,
+                unittest.mock.ANY,
+                unittest.mock.ANY,
+                response.headers['correlation-id']
+            )
+
+    def test_gte_500_logs_error(self):
+        for status in (500, 504):
+            response = self.fetch('/status/{}'.format(status))
+
+            self.access_logger.error.assert_any_call(
+                "%d %s %.2fms {CID %s}",
+                status,
+                unittest.mock.ANY,
+                unittest.mock.ANY,
+                response.headers['correlation-id']
+            )
+
+    def test_uses_correlation_id_from_header_if_missing_from_handler(self):
+        correlation_id = uuid.uuid4().hex
+
+        self.fetch('/status/no-correlation/200',
+                   headers={'Correlation-Id': correlation_id})
+
+        self.access_logger.info.assert_any_call(
+            "%d %s %.2fms {CID %s}",
+            200,
+            unittest.mock.ANY,
+            unittest.mock.ANY,
+            correlation_id
+        )
+
+    def test_correlation_id_is_none_if_missing_from_handler_and_header(self):
+        self.fetch('/status/no-correlation/200')
+
+        self.access_logger.info.assert_any_call(
+            "%d %s %.2fms {CID %s}",
+            200,
+            unittest.mock.ANY,
+            unittest.mock.ANY,
+            None
+        )