reviews and tests

logging/google/cloud/logging/_shutdown.py

@@ -15,7 +15,6 @@
 """Module for code related to shutdown logging."""
 
 import functools
-import httplib2
 import os
 import signal
 import sys
@@ -24,68 +23,81 @@
 from google.cloud.logging.environment_vars import _APPENGINE_FLEXIBLE_ENV_VM
 from google.cloud.logging.environment_vars import _APPENGINE_FLEXIBLE_ENV_FLEX
 
-_METADATA_SERVICE_ADDR = '169.254.169.254'
-_APPENGINE_SERVICE = 'appengine.googleapis.com'
 
-
-def _fetch_metadata(metdata_path):
-    url = 'http://{}/computeMetadata/v1/{}'.format(_METADATA_SERVICE_ADDR,
-                                                   metdata_path)
-    h = httplib2.Http()
-    headers = {'Metadata-Flavor': 'Google'}
-    resp, content = h.request(url, 'GET')
-
-    return content
+# Maximum size in bytes to send to Stackdriver Logging in one entry
+MAX_PAYLOAD_SIZE = 1024 * 100
 
 
 def _get_gae_instance():
-    return _fetch_metadata('instance/attributes/gae_backend_instance')
+    """Returns the App Engine Flexible instance."""
+    return os.getenv('GAE_INSTANCE')
 
 
-def _get_gae_backend():
-    return _fetch_metadata('instance/attributes/gae_backend_name')
+def _get_gae_service():
+    """Returns the App Engine Flexible service."""
+    return os.getenv('GAE_SERVICE')
 
 
 def _get_gae_version():
-    return _fetch_metadata('instance/attributes/gae_backend_version')
+    """Returns the App Engine Flexible version."""
+    return os.getenv('GAE_VERSION')
+
+
+def _split_entry(payload):
+    """Splits payload into lists of maximum 100Kb.
+
+    Stackdriver Logging payloads are a maximum of 100Kb.
+    """
+    return [payload[i:i + MAX_PAYLOAD_SIZE]
+            for i in range(0, len(payload), MAX_PAYLOAD_SIZE)]
 
 
 def _write_stacktrace_log(client, traces):
+    """Writes the trace logs to the appropriate GAE resource in Stackdriver.
+
+    :type client: `google.cloud.logging.Client`
+    :param client: Stackdriver logging client.
+
+    :type traces: str
+    :param traces: String containing the stacktrace info to write to
+                   Stackdriver logging.
+    """
     gae_version = _get_gae_version()
-    gae_backend = _get_gae_backend()
+    gae_service = _get_gae_service()
     gae_instance = _get_gae_instance()
 
     text_payload = '{}\nThread traces\n{}'.format(gae_instance, traces)
-    logger_name = 'projects/{}/logs/'
-    'appengine.googleapis.com%2Fapp.shutdown'.format(client.project)
+    logger_name = 'projects/{}/logs/appengine.googleapis.com%2Fapp.shutdown'.format(client.project)
 
     resource = {'type': 'gae_app', 'labels': {'project_id': client.project,
                                               'version_id': gae_version,
-                                              'module_id': gae_backend}}
+                                              'module_id': gae_service}}
 
     labels = {'appengine.googleapis.com/version_id': gae_version,
               'compute.googleapis.com/resource_type': 'instance',
               'appengine.googleapis.com/instance_name': gae_instance,
-              'appengine.googleapis.com / module_id': gae_backend, }
-    entry = {'text_payload': text_payload}
+              'appengine.googleapis.com / module_id': gae_service, }
 
-    entries = [entry]
+    split_payloads = _split_entry(bytes(text_payload))
+    entries = [{'text_payload': payload} for payload in split_payloads]
 
+    print "SENDINT ENTRY %s logger  %s resource %s " % (entries, logger_name,
+                                                        resource )
     client.logging_api.write_entries(
         entries, logger_name=logger_name, resource=resource, labels=labels)
 
 
 def _is_on_appengine():
+    """Returns True if the environment is detected as App Engine flexible."""
     return (os.getenv(_APPENGINE_FLEXIBLE_ENV_VM) or os.getenv(
         _APPENGINE_FLEXIBLE_ENV_FLEX))
 
 
 def _report_stacktraces(client, signal, frame):
-    """
-    Reports the stacktraces of all active threads to Stackdriver Logging.
+    """Reports the stacktraces of all active threads to Stackdriver Logging.
 
     :type client: `google.cloud.logging.Client`
-    :param client: Stackdriver loggingclient.
+    :param client: Stackdriver logging client.
 
     :type signal: int
     :param signal: Signal number.
@@ -94,13 +106,19 @@ def _report_stacktraces(client, signal, frame):
     :param frame: The current stack frame.
     """
     traces = ''
+    print "RUNNING SIGNAL HANDLER!"
     for threadId, stack in sys._current_frames().items():
         traces += '\n# ThreadID: {}'.format(threadId)
         for filename, lineno, name, line in traceback.extract_stack(stack):
             traces += 'File: {}, line {}, in {}'.format(
                 filename, lineno, name)
+
     _write_stacktrace_log(client, traces)
 
+def fml(signal, frame):
+    print "MY LIFE IS F'ed"
+
+
 
 def setup_shutdown_stacktrace_reporting(client):
     """Installs a SIGTERM handler to log stack traces to  Stackdriver.
@@ -109,6 +127,7 @@ def setup_shutdown_stacktrace_reporting(client):
     :param client: Stackdriver logging client.
     """
     if not _is_on_appengine():
-        raise Exception('Shutdown reporting is only supported on App Engine '
-                        'flexible environment.')
+        raise RuntimeError('Shutdown reporting is only supported on App '
+                           'Engine flexible environment.')
+    print "INSTALLIGN SIGNAL HANDLER"
     signal.signal(signal.SIGTERM, functools.partial(_report_stacktraces, client))

logging/unit_tests/test__shutdown.py

@@ -0,0 +1,86 @@
+# Copyright 2017 Google Inc.
+#
+# 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
+
+import mock
+
+
+class Test_setup_shutdown_stracktrace_reporting(unittest.TestCase):
+
+    def _call_fut(self, request):
+        from google.cloud.logging._shutdown import (
+            setup_shutdown_stacktrace_reporting)
+        return setup_shutdown_stacktrace_reporting(request)
+
+    def test_setup_shutdown_stacktrace_reporting_no_gae(self):
+        with self.assertRaises(RuntimeError):
+            self._call_fut(mock.Mock())
+
+    def test_setup_shutdown_stacktrace_reporting(self):
+        from google.cloud.logging.environment_vars import (
+            _APPENGINE_FLEXIBLE_ENV_VM)
+        from google.cloud._testing import _Monkey
+        import os
+
+        signal_patch = mock.patch('google.cloud.logging._shutdown.signal')
+
+        with _Monkey(os, environ={_APPENGINE_FLEXIBLE_ENV_VM: 'True'}):
+            with signal_patch as signal_mock:
+                self._call_fut(mock.Mock())
+                self.assertTrue(signal_mock.signal.called)
+
+
+class Test_write_stackrace_log(unittest.TestCase):
+
+    def _call_fut(self, client, traces):
+        from google.cloud.logging._shutdown import (
+            _write_stacktrace_log
+        )
+        return _write_stacktrace_log(client, traces)
+
+
+    def test_write_stracktrace_log(self):
+        from google.cloud._testing import _Monkey
+        import os
+
+        mock_client = mock.Mock()
+
+        trace = 'a simple stack trace'
+        with _Monkey(os, environ={'GAE_INSTANCE': 'myversion'}):
+            self._call_fut(mock_client, trace)
+            called = mock_client.logging_api.write_entries.call_args
+
+        expected_payload = 'myversion\nThread traces\n{}'.format(trace)
+        self.assertEqual(called[0][0], [{'text_payload': expected_payload}])
+
+
+class Test_report_stacktraces(unittest.TestCase):
+
+    def _call_fut(self, client, signal, frame):
+        from google.cloud.logging._shutdown import (
+            _report_stacktraces
+        )
+        return _report_stacktraces(client, signal, frame)
+
+    def test_report_stacktraces(self):
+        patch = mock.patch(
+            'google.cloud.logging._shutdown._write_stacktrace_log')
+        with patch as write_log_mock:
+            self._call_fut(mock.Mock(), mock.Mock() ,mock.Mock())
+
+            traces = write_log_mock.call_args[0][1]
+        self.assertIn('test__shutdown', traces)
+
+

logging/unit_tests/test_client.py

@@ -632,6 +632,22 @@ def test_setup_logging(self):
 
         setup_logging.assert_called()
 
+    def test_setup_shutdown_stacktrace_reporting(self):
+        import mock
+
+        credentials = _make_credentials()
+        shutdown_patch = mock.patch(
+            'google.cloud.logging.client.setup_shutdown_stacktrace_reporting')
+        with shutdown_patch as shutdown_mock:
+            client = self._make_one(project=self.PROJECT,
+                                    credentials=credentials,
+                                    use_gax=False)
+            client.enable_shutdown_logging(thread_dump=False)
+            shutdown_mock.assert_not_called()
+
+            client.enable_shutdown_logging(thread_dump=True)
+            shutdown_mock.assert_called()
+
 
 class _Connection(object):