1#!/usr/bin/python
2# Copyright 2016 Google Inc. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Unittest for compute_auth.py module."""
17
18import sys
19
20from google_compute_engine.test_compat import mock
21from google_compute_engine.test_compat import unittest
22
23if sys.version_info < (3, 0):
24  from google_compute_engine.boto import compute_auth
25
26
27@unittest.skipIf(sys.version_info > (3, 0), 'Skipping for python3.')
28class ComputeAuthTest(unittest.TestCase):
29
30  def setUp(self):
31    self.metadata_key = 'instance/service-accounts'
32    self.service_account = 'service_account'
33    self.mock_config = mock.Mock()
34    self.mock_config.get.return_value = self.service_account
35    self.mock_provider = mock.Mock()
36    self.mock_provider.name = 'google'
37
38  @mock.patch('google_compute_engine.boto.compute_auth.metadata_watcher')
39  @mock.patch('google_compute_engine.boto.compute_auth.logger')
40  def testCreateConfig(self, mock_logger, mock_watcher):
41    scopes = list(compute_auth.GS_SCOPES)[1:2]
42    service_accounts = {self.service_account: {'scopes': scopes}}
43    mock_watcher.GetMetadata.return_value = service_accounts
44    mock_watcher.MetadataWatcher.return_value = mock_watcher
45    mocks = mock.Mock()
46    mocks.attach_mock(mock_watcher, 'watcher')
47    mocks.attach_mock(mock_logger, 'logger')
48    mocks.attach_mock(self.mock_config, 'config')
49    mock_logger_instance = mock.Mock()
50    mock_logger.Logger.return_value = mock_logger_instance
51
52    mock_compute_auth = compute_auth.ComputeAuth(
53        None, self.mock_config, self.mock_provider)
54    expected_calls = [
55        mock.call.logger.Logger(name=mock.ANY),
56        mock.call.watcher.MetadataWatcher(logger=mock_logger_instance),
57        mock.call.config.get('GoogleCompute', 'service_account', ''),
58        mock.call.watcher.GetMetadata(metadata_key=self.metadata_key)
59    ]
60    self.assertEqual(mocks.mock_calls, expected_calls)
61    self.assertEqual(mock_compute_auth.scopes, scopes)
62
63  @mock.patch('google_compute_engine.boto.compute_auth.metadata_watcher')
64  def testCreateConfigNoScopes(self, mock_watcher):
65    mock_watcher.GetMetadata.return_value = {}
66    mock_watcher.MetadataWatcher.return_value = mock_watcher
67
68    with self.assertRaises(compute_auth.auth_handler.NotReadyToAuthenticate):
69      compute_auth.ComputeAuth(None, self.mock_config, self.mock_provider)
70
71  def testCreateConfigNoServiceAccount(self):
72    self.mock_config.get.return_value = None
73
74    with self.assertRaises(compute_auth.auth_handler.NotReadyToAuthenticate):
75      compute_auth.ComputeAuth(None, self.mock_config, self.mock_provider)
76
77  @mock.patch('google_compute_engine.boto.compute_auth.metadata_watcher')
78  def testGetAccessToken(self, mock_watcher):
79    mock_auth = mock.create_autospec(compute_auth.ComputeAuth)
80    mock_auth.watcher = mock_watcher
81    mock_auth.metadata_key = self.metadata_key
82    mock_auth.service_account = self.service_account
83    mock_watcher.GetMetadata.side_effect = [
84        {self.service_account: {'token': {'access_token': 'test'}}},
85        {},
86    ]
87
88    self.assertEqual(
89        compute_auth.ComputeAuth._GetAccessToken(mock_auth), 'test')
90    self.assertEqual(
91        compute_auth.ComputeAuth._GetAccessToken(mock_auth), None)
92    expected_calls = [mock.call(metadata_key=self.metadata_key)] * 2
93    self.assertEqual(mock_watcher.GetMetadata.mock_calls, expected_calls)
94
95  @mock.patch('google_compute_engine.boto.compute_auth.metadata_watcher')
96  def testAddAuth(self, mock_watcher):
97    mock_auth = mock.create_autospec(compute_auth.ComputeAuth)
98    mock_auth._GetAccessToken.return_value = 'token'
99    mock_request = mock.Mock()
100    mock_request.headers = {}
101
102    compute_auth.ComputeAuth.add_auth(mock_auth, mock_request)
103    self.assertEqual(mock_request.headers['Authorization'], 'OAuth token')
104
105
106if __name__ == '__main__':
107  unittest.main()
108