sdks/python/apache_beam/internal/gcp/auth_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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 logging
import unittest

import mock

from apache_beam.internal.gcp import auth
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions

try:
  import google.auth as gauth
except ImportError:
  gauth = None  # type: ignore


class MockLoggingHandler(logging.Handler):
  """Mock logging handler to check for expected logs."""
  def __init__(self, *args, **kwargs):
    self.reset()
    logging.Handler.__init__(self, *args, **kwargs)

  def emit(self, record):
    self.messages[record.levelname.lower()].append(record.getMessage())

  def reset(self):
    self.messages = {
        'debug': [],
        'info': [],
        'warning': [],
        'error': [],
        'critical': [],
    }


@unittest.skipIf(gauth is None, 'Google Auth dependencies are not installed')
class AuthTest(unittest.TestCase):
  @mock.patch('google.auth.default')
  def test_auth_with_retrys(self, unused_mock_arg):
    pipeline_options = PipelineOptions()
    pipeline_options.view_as(
        GoogleCloudOptions).impersonate_service_account = False

    credentials = ('creds', 1)

    self.is_called = False

    def side_effect(scopes=None):
      if self.is_called:
        return credentials
      else:
        self.is_called = True
        raise IOError('Failed')

    google_auth_mock = mock.MagicMock()
    gauth.default = google_auth_mock
    google_auth_mock.side_effect = side_effect

    # _Credentials caches the actual credentials.
    # This resets it for idempotent tests.
    if auth._Credentials._credentials_init:
      auth._Credentials._credentials_init = False
      auth._Credentials._credentials = None

    returned_credentials = auth.get_service_credentials(pipeline_options)

    # _Credentials caches the actual credentials.
    # This resets it for idempotent tests.
    if auth._Credentials._credentials_init:
      auth._Credentials._credentials_init = False
      auth._Credentials._credentials = None

    self.assertEqual('creds', returned_credentials._google_auth_credentials)

  @mock.patch(
      'apache_beam.internal.gcp.auth._Credentials._get_credentials_with_retrys')
  def test_auth_with_retrys_always_fail(self, unused_mock_arg):
    pipeline_options = PipelineOptions()
    pipeline_options.view_as(
        GoogleCloudOptions).impersonate_service_account = False

    loggerHandler = MockLoggingHandler()

    auth._LOGGER.addHandler(loggerHandler)

    #Remove call to retrying method, as otherwise test takes ~10 minutes to run
    def raise_(scopes=None):
      raise IOError('Failed')

    retry_auth_mock = mock.MagicMock()
    auth._Credentials._get_credentials_with_retrys = retry_auth_mock
    retry_auth_mock.side_effect = raise_

    # _Credentials caches the actual credentials.
    # This resets it for idempotent tests.
    if auth._Credentials._credentials_init:
      auth._Credentials._credentials_init = False
      auth._Credentials._credentials = None

    returned_credentials = auth.get_service_credentials(pipeline_options)

    self.assertEqual(None, returned_credentials)
    self.assertEqual([
        'Unable to find default credentials to use: Failed\n'
        'Connecting anonymously. This is expected if no credentials are '
        'needed to access GCP resources.'
    ],
                     loggerHandler.messages.get('warning'))

    # _Credentials caches the actual credentials.
    # This resets it for idempotent tests.
    if auth._Credentials._credentials_init:
      auth._Credentials._credentials_init = False
      auth._Credentials._credentials = None

    auth._LOGGER.removeHandler(loggerHandler)


if __name__ == '__main__':
  unittest.main()