Add unit test for flask and some refactor

logging/google/cloud/logging/handlers/_helpers.py

@@ -27,6 +27,8 @@
 _FLASK_TRACE_HEADER = 'X_CLOUD_TRACE_CONTEXT'
 _DJANGO_TRACE_HEADER = 'HTTP_X_CLOUD_TRACE_CONTEXT'
 
+_EMPTY_TRACE_ID = 'None'
+
 
 def format_stackdriver_json(record, message):
     """Helper to format a LogRecord in in Stackdriver fluentd format.
@@ -56,12 +58,12 @@ def get_trace_id_from_flask():
     :return: Trace_id in HTTP request headers.
     """
     if not flask or not flask.request:
-        return None
+        return _EMPTY_TRACE_ID
 
     header = flask.request.headers.get(_FLASK_TRACE_HEADER)
 
     if not header:
-        return None
+        return _EMPTY_TRACE_ID
 
     trace_id = header.split('/')[0]
 
@@ -78,12 +80,12 @@ def get_trace_id_from_django():
     request = request_middleware.get_request()
 
     if not request:
-        return None
+        return _EMPTY_TRACE_ID
 
     try:
         header = request.META[_DJANGO_TRACE_HEADER]
     except KeyError:
-        return None
+        return _EMPTY_TRACE_ID
 
     trace_id = header.split('/')[0]
 
@@ -100,7 +102,7 @@ def get_trace_id():
 
     for checker in checkers:
         trace_id = checker()
-        if trace_id is not None:
+        if trace_id is not _EMPTY_TRACE_ID:
             return trace_id
 
     return trace_id

logging/google/cloud/logging/handlers/app_engine.py

@@ -77,8 +77,6 @@ def get_gae_labels(self):
         :returns: Labels for GAE app.
         """
         trace_id = get_trace_id()
-        if trace_id is None:
-            trace_id = 'unknown'
         gae_labels = {
             'appengine.googleapis.com/trace_id': trace_id,
         }

logging/nox.py

@@ -31,7 +31,7 @@ def unit_tests(session, python_version):
     session.interpreter = 'python{}'.format(python_version)
 
     # Install all test dependencies, then install this package in-place.
-    session.install('mock', 'pytest', 'pytest-cov', *LOCAL_DEPS)
+    session.install('mock', 'pytest', 'pytest-cov', 'flask', *LOCAL_DEPS)
     session.install('-e', '.')
 
     # Run py.test against the unit tests.

logging/tests/unit/handlers/test__helpers.py

@@ -0,0 +1,57 @@
+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+
+class TestFlaskTrace(unittest.TestCase):
+
+    def create_app(self):
+        import flask
+
+        app = flask.Flask(__name__)
+
+        @app.route('/')
+        def index():
+            return 'test flask trace'
+
+        return app
+
+    def setUp(self):
+        self.app = self.create_app()
+
+    def test_trace_id_no_context_header(self):
+        from google.cloud.logging.handlers._helpers import get_trace_id_from_flask
+        from google.cloud.logging.handlers._helpers import _EMPTY_TRACE_ID
+
+        with self.app.test_request_context(
+                path='/',
+                headers={}):
+            trace_id = get_trace_id_from_flask()
+
+        self.assertEqual(trace_id, _EMPTY_TRACE_ID)
+
+    def test_trace_id_valid_context_header(self):
+        from google.cloud.logging.handlers._helpers import get_trace_id_from_flask
+
+        FLASK_TRACE_HEADER = 'X_CLOUD_TRACE_CONTEXT'
+        FLASK_TRACE_ID = 'testtraceid/testspanid'
+
+        with self.app.test_request_context(
+                path='/',
+                headers={FLASK_TRACE_HEADER:FLASK_TRACE_ID}):
+            trace_id = get_trace_id_from_flask()
+
+        EXPECTED_TRACE_ID = 'testtraceid'
+        self.assertEqual(trace_id, EXPECTED_TRACE_ID)

logging/tests/unit/handlers/test_app_engine.py

@@ -32,6 +32,7 @@ def test_constructor(self):
         from google.cloud.logging.handlers.app_engine import _GAE_PROJECT_ENV
         from google.cloud.logging.handlers.app_engine import _GAE_SERVICE_ENV
         from google.cloud.logging.handlers.app_engine import _GAE_VERSION_ENV
+        from google.cloud.logging.handlers._helpers import _EMPTY_TRACE_ID
 
         TRACE_ID_LABEL = 'appengine.googleapis.com/trace_id'
 
@@ -46,7 +47,7 @@ def test_constructor(self):
         self.assertEqual(handler.resource.labels['project_id'], 'test_project')
         self.assertEqual(handler.resource.labels['module_id'], 'test_service')
         self.assertEqual(handler.resource.labels['version_id'], 'test_version')
-        self.assertEqual(handler.labels[TRACE_ID_LABEL], 'unknown')
+        self.assertEqual(handler.labels[TRACE_ID_LABEL], _EMPTY_TRACE_ID)
 
     def test_emit(self):
         import mock