Add Cloud Logging initialization based on ENV VARs and unit test.

If CHROMITE_CLOUD_LOGGING and GOOGLE_APPLICATION_CREDENTIALS are not
both set then no cros_logging is not affected for chromite clients.

I have manually verified in my own chroot that if I set
CHROMITE_CLOUD_LOGGING=1 and GOOGLE_APPLICATION_CREDENTIALS=<local.json>
where the local json file for a google-cloud-logging-enabled service account
I created in GCP then I do end up with log statements from a script run
in the chroot going to this GCP project that are visible in Google cloud
logging.

See https://cloud.google.com/docs/authentication/getting-started#cloud-console

Next step: enable these variables in staging-only, with
GOOGLE_APPLICTION_CREDENTIALS referring to a bot-based json file for the
chromebot account.  Test staging setting in CQ or with led
(go/luci-how-to-led).

BUG=chromium:1128411
TEST=manual test, pytest

Change-Id: I59d11b927fbb89b641aedb43d944918a343fccf6
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2472537
Tested-by: Michael Mortensen <mmortensen@google.com>
Reviewed-by: Alex Klein <saklein@chromium.org>
Commit-Queue: Michael Mortensen <mmortensen@google.com>
diff --git a/lib/cros_logging.py b/lib/cros_logging.py
index 38855b7..a713aac 100644
--- a/lib/cros_logging.py
+++ b/lib/cros_logging.py
@@ -39,6 +39,7 @@
 
 from __future__ import print_function
 
+import os
 import sys
 # pylint: disable=unused-wildcard-import, wildcard-import
 from logging import *
@@ -53,7 +54,6 @@
 # Import as private to avoid polluting module namespace.
 from chromite.lib import buildbot_annotations as _annotations
 
-
 # Remove deprecated APIs to force use of new ones.
 del WARN
 del warn
@@ -64,6 +64,50 @@
 addLevelName(NOTICE, 'NOTICE')
 
 
+def _SetupCloudLogging():
+  """If appropriate environment variables are set, enable cloud logging.
+
+  Cloud logging is only enabled when the environment has
+   CHROMITE_CLOUD_LOGGING=1 and GOOGLE_APPLICATION_CREDENTIALS=<local.json>.
+  If these are set, then cloud logging is enable, see
+  https://cloud.google.com/docs/authentication/getting-started#cloud-console
+  """
+  try:
+    import google.cloud.logging as cloud_logging  # pylint: disable=import-error,no-name-in-module
+  except ImportError as e:
+    # TODO(mmortensen): Change to python3's ModuleNotFoundError when this
+    # code is only used by python3. Beware though with branches and bisection
+    # this could need to be ImportError for a long time. ImportError is the
+    # parent class of ModuleNotFoundError and works on both python2 and python3.
+    log(NOTICE, 'Could not import google.cloud.logging %s', e)
+    return
+
+  client = cloud_logging.Client()
+  # Retrieves a Cloud Logging handler based on the environment
+  # you're running in and integrates the handler with the
+  # Python logging module. By default this captures all logs
+  # at INFO level and higher
+  client.get_default_handler()
+  client.setup_logging()
+
+
+def _CloudLoggingEnvVariablesAreDefined():
+  """Check for cloud-logging ENV variables."""
+  cloud_logging_env_value = os.environ.get('CHROMITE_CLOUD_LOGGING')
+  google_app_creds_env_value = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
+  # If both variables are set, log their values and return True.
+  if cloud_logging_env_value == '1' and google_app_creds_env_value:
+    log(NOTICE, 'CHROMITE_CLOUD_LOGGING is ', cloud_logging_env_value)
+    log(NOTICE,
+        'GOOGLE_APPLICATION_CREDENTIALS is ', google_app_creds_env_value)
+    return True
+  return False
+
+
+if _CloudLoggingEnvVariablesAreDefined():
+  _SetupCloudLogging()
+
+
 # Notice implementation.
 def notice(message, *args, **kwargs):
   """Log 'msg % args' with severity 'NOTICE'."""
diff --git a/lib/cros_logging_unittest.py b/lib/cros_logging_unittest.py
index b18ac7e..228675a 100644
--- a/lib/cros_logging_unittest.py
+++ b/lib/cros_logging_unittest.py
@@ -7,12 +7,78 @@
 
 from __future__ import print_function
 
+import os
 import sys
 
 from chromite.lib import cros_logging as logging
 from chromite.lib import cros_test_lib
 
 
+class CrosCloudloggingTest(cros_test_lib.MockOutputTestCase):
+  """Test google-cloud-logging interacts with logging as expected."""
+
+  # pylint: disable=protected-access
+
+  def setUp(self):
+    self.logger = logging.getLogger()
+    sh = logging.StreamHandler(sys.stdout)
+    self.logger.addHandler(sh)
+    # pylint: disable=protected-access
+    logging._buildbot_markers_enabled = False
+    try:
+      import google.cloud.logging  # pylint: disable=import-error,no-name-in-module
+    except ImportError:
+      self.cloud_logging_import_error = True
+      return
+    self.cloud_logging_import_error = False
+    self.client_mock = self.PatchObject(google.cloud.logging, 'Client')
+
+  def testSetupCloudLogging(self):
+    if self.cloud_logging_import_error:
+      return
+    ### client_mock = self.PatchObject(google.cloud.logging, 'Client')
+    # Invoke the code that calls logging when env vars are set.
+    logging._SetupCloudLogging()
+
+    # Verify that google.cloud.logging.Client() was executed.
+    self.client_mock.assert_called_once()
+
+  def testCloudLoggingEnvVariablesAreDefined_notSet(self):
+    if self.cloud_logging_import_error:
+      return
+    with self.OutputCapturer() as output:
+      cloud_env_defined = logging._CloudLoggingEnvVariablesAreDefined()
+    # If both are not set, there should be no output.
+    self.assertEqual(output.GetStdout(), '')
+    self.assertEqual(output.GetStderr(), '')
+    self.assertFalse(cloud_env_defined)
+
+  def testCloudLoggingEnvVariablesAreDefined_envSet(self):
+    if self.cloud_logging_import_error:
+      return
+    # Set the env vars
+    os.environ['CHROMITE_CLOUD_LOGGING'] = '1'
+    os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '/some/path/to/creds.json'
+    with self.OutputCapturer() as output:
+      cloud_env_defined = logging._CloudLoggingEnvVariablesAreDefined()
+    # Verify that both variables are logged.
+    self.assertTrue(cloud_env_defined)
+    self.assertIn('CHROMITE_CLOUD_LOGGING', output.GetStderr())
+    self.assertIn('GOOGLE_APPLICATION_CREDENTIALS', output.GetStderr())
+
+  def testCloudLoggingEnvVariablesAreDefined_noAllEnvSet(self):
+    if self.cloud_logging_import_error:
+      return
+    # Set the env vars
+    os.environ['CHROMITE_CLOUD_LOGGING'] = '1'
+    with self.OutputCapturer() as output:
+      cloud_env_defined = logging._CloudLoggingEnvVariablesAreDefined()
+    # If both are not set, there should be no output.
+    self.assertEqual(output.GetStdout(), '')
+    self.assertEqual(output.GetStderr(), '')
+    self.assertFalse(cloud_env_defined)
+
+
 class CrosloggingTest(cros_test_lib.OutputTestCase):
   """Test logging works as expected."""