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 metadata_watcher.py module."""
17
18import os
19
20from google_compute_engine import metadata_watcher
21from google_compute_engine.test_compat import mock
22from google_compute_engine.test_compat import unittest
23
24
25class MetadataWatcherTest(unittest.TestCase):
26
27  def setUp(self):
28    self.mock_logger = mock.Mock()
29    self.timeout = 60
30    self.url = 'http://metadata.google.internal/computeMetadata/v1'
31    self.params = {
32        'alt': 'json',
33        'last_etag': 0,
34        'recursive': True,
35        'timeout_sec': self.timeout,
36        'wait_for_change': True,
37    }
38    self.mock_watcher = metadata_watcher.MetadataWatcher(
39        logger=self.mock_logger, timeout=self.timeout)
40
41  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.build_opener')
42  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.ProxyHandler')
43  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.Request')
44  def testGetMetadataRequest(self, mock_request, mock_proxy, mock_opener):
45    mock_open = mock.Mock()
46    mock_handler = mock.Mock()
47    mock_response = mock.Mock()
48    mocks = mock.Mock()
49    mocks.attach_mock(mock_request, 'request')
50    mocks.attach_mock(mock_proxy, 'proxy')
51    mocks.attach_mock(mock_handler, 'handler')
52    mocks.attach_mock(mock_opener, 'opener')
53    mocks.attach_mock(mock_open, 'open')
54    mocks.attach_mock(mock_response, 'response')
55    mock_request.return_value = mock_request
56    mock_proxy.return_value = mock_handler
57    mock_opener.return_value = mock_open
58    mock_response.getcode.return_value = metadata_watcher.httpclient.OK
59    mock_open.open.return_value = mock_response
60    params = {'hello': 'world'}
61    request_url = '%s?hello=world' % self.url
62    headers = {'Metadata-Flavor': 'Google'}
63    timeout = self.timeout * 1.1
64
65    self.mock_watcher._GetMetadataRequest(self.url, params=params)
66    expected_calls = [
67        mock.call.request(request_url, headers=headers),
68        mock.call.proxy({}),
69        mock.call.opener(mock_handler),
70        mock.call.open.open(mock_request, timeout=timeout),
71        mock.call.response.getcode(),
72    ]
73    self.assertEqual(mocks.mock_calls, expected_calls)
74
75  @mock.patch('google_compute_engine.metadata_watcher.time')
76  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.build_opener')
77  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.ProxyHandler')
78  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.Request')
79  def testGetMetadataRequestRetry(
80      self, mock_request, mock_proxy, mock_opener, mock_time):
81    mock_open = mock.Mock()
82    mock_handler = mock.Mock()
83    mocks = mock.Mock()
84    mocks.attach_mock(mock_request, 'request')
85    mocks.attach_mock(mock_proxy, 'proxy')
86    mocks.attach_mock(mock_handler, 'handler')
87    mocks.attach_mock(mock_opener, 'opener')
88    mocks.attach_mock(mock_open, 'open')
89    mocks.attach_mock(mock_time, 'time')
90    mock_request.return_value = mock_request
91    mock_proxy.return_value = mock_handler
92    mock_opener.return_value = mock_open
93
94    mock_unavailable = mock.Mock()
95    mock_unavailable.getcode.return_value = (
96        metadata_watcher.httpclient.SERVICE_UNAVAILABLE)
97    mock_timeout = metadata_watcher.socket.timeout('Test timeout')
98    mock_success = mock.Mock()
99    mock_success.getcode.return_value = metadata_watcher.httpclient.OK
100
101    # Retry after a service unavailable error response.
102    mock_open.open.side_effect = [
103        metadata_watcher.StatusException(mock_unavailable),
104        mock_timeout,
105        mock_success,
106    ]
107    request_url = '%s?' % self.url
108    headers = {'Metadata-Flavor': 'Google'}
109    timeout = self.timeout * 1.1
110
111    self.mock_watcher._GetMetadataRequest(self.url)
112    expected_calls = [
113        mock.call.request(request_url, headers=headers),
114        mock.call.proxy({}),
115        mock.call.opener(mock_handler),
116        mock.call.open.open(mock_request, timeout=timeout),
117        mock.call.time.sleep(mock.ANY),
118        mock.call.request(request_url, headers=headers),
119        mock.call.proxy({}),
120        mock.call.opener(mock_handler),
121        mock.call.open.open(mock_request, timeout=timeout),
122        mock.call.time.sleep(mock.ANY),
123        mock.call.request(request_url, headers=headers),
124        mock.call.proxy({}),
125        mock.call.opener(mock_handler),
126        mock.call.open.open(mock_request, timeout=timeout),
127    ]
128    self.assertEqual(mocks.mock_calls, expected_calls)
129
130  @mock.patch('google_compute_engine.metadata_watcher.time')
131  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.build_opener')
132  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.ProxyHandler')
133  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.Request')
134  def testGetMetadataRequestHttpException(
135      self, mock_request, mock_proxy, mock_opener, mock_time):
136    mock_open = mock.Mock()
137    mock_handler = mock.Mock()
138    mock_response = mock.Mock()
139    mock_request.return_value = mock_request
140    mock_proxy.return_value = mock_handler
141    mock_opener.return_value = mock_open
142    mock_response.getcode.return_value = metadata_watcher.httpclient.NOT_FOUND
143    mock_open.open.side_effect = metadata_watcher.StatusException(mock_response)
144
145    with self.assertRaises(metadata_watcher.StatusException):
146      self.mock_watcher._GetMetadataRequest(self.url)
147    self.assertEqual(mock_request.call_count, 1)
148    self.assertEqual(mock_proxy.call_count, 1)
149    self.assertEqual(mock_opener.call_count, 1)
150    self.assertEqual(mock_open.open.call_count, 1)
151    self.assertEqual(mock_response.getcode.call_count, 1)
152
153  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.build_opener')
154  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.ProxyHandler')
155  @mock.patch('google_compute_engine.metadata_watcher.urlrequest.Request')
156  def testGetMetadataRequestException(
157      self, mock_request, mock_proxy, mock_opener):
158    mock_open = mock.Mock()
159    mock_handler = mock.Mock()
160    mock_response = mock.Mock()
161    mock_request.return_value = mock_request
162    mock_proxy.return_value = mock_handler
163    mock_opener.return_value = mock_open
164    mock_response.getcode.return_value = metadata_watcher.httpclient.NOT_FOUND
165    mock_open.open.side_effect = mock_response
166
167    with self.assertRaises(metadata_watcher.StatusException):
168      self.mock_watcher._GetMetadataRequest(self.url)
169    self.assertEqual(mock_request.call_count, 1)
170    self.assertEqual(mock_proxy.call_count, 1)
171    self.assertEqual(mock_opener.call_count, 1)
172    self.assertEqual(mock_open.open.call_count, 1)
173
174  def testUpdateEtag(self):
175    mock_response = mock.Mock()
176    mock_response.headers = {'etag': 1}
177    self.assertEqual(self.mock_watcher.etag, 0)
178
179    # Update the etag if the etag is set.
180    self.assertTrue(self.mock_watcher._UpdateEtag(mock_response))
181    self.assertEqual(self.mock_watcher.etag, 1)
182
183    # Do not update the etag if the etag is unchanged.
184    self.assertFalse(self.mock_watcher._UpdateEtag(mock_response))
185    self.assertEqual(self.mock_watcher.etag, 1)
186
187    # Do not update the etag if the etag is not set.
188    mock_response.headers = {}
189    self.assertFalse(self.mock_watcher._UpdateEtag(mock_response))
190    self.assertEqual(self.mock_watcher.etag, 1)
191
192  def testGetMetadataUpdate(self):
193    mock_response = mock.Mock()
194    mock_response.return_value = mock_response
195    mock_response.headers = {'etag': 1}
196    mock_response.read.return_value = bytes(b'{}')
197    self.mock_watcher._GetMetadataRequest = mock_response
198    request_url = os.path.join(self.url, '')
199
200    self.assertEqual(self.mock_watcher._GetMetadataUpdate(), {})
201    self.assertEqual(self.mock_watcher.etag, 1)
202    mock_response.assert_called_once_with(
203        request_url, params=self.params, timeout=None)
204
205  def testGetMetadataUpdateArgs(self):
206    mock_response = mock.Mock()
207    mock_response.return_value = mock_response
208    mock_response.headers = {'etag': 0}
209    mock_response.read.return_value = bytes(b'{}')
210    self.mock_watcher._GetMetadataRequest = mock_response
211    metadata_key = 'instance/id'
212    self.params['recursive'] = False
213    self.params['wait_for_change'] = False
214    request_url = os.path.join(self.url, metadata_key)
215
216    self.mock_watcher._GetMetadataUpdate(
217        metadata_key=metadata_key, recursive=False, wait=False, timeout=60)
218    self.assertEqual(self.mock_watcher.etag, 0)
219    mock_response.assert_called_once_with(
220        request_url, params=self.params, timeout=60)
221
222  def testGetMetadataUpdateWait(self):
223    self.params['last_etag'] = 1
224    self.mock_watcher.etag = 1
225    mock_unchanged = mock.Mock()
226    mock_unchanged.headers = {'etag': 1}
227    mock_unchanged.read.return_value = bytes(b'{}')
228    mock_changed = mock.Mock()
229    mock_changed.headers = {'etag': 2}
230    mock_changed.read.return_value = bytes(b'{}')
231    mock_response = mock.Mock()
232    mock_response.side_effect = [mock_unchanged, mock_unchanged, mock_changed]
233    self.mock_watcher._GetMetadataRequest = mock_response
234    request_url = os.path.join(self.url, '')
235
236    self.mock_watcher._GetMetadataUpdate()
237    self.assertEqual(self.mock_watcher.etag, 2)
238    expected_calls = [
239        mock.call(request_url, params=self.params, timeout=None),
240    ] * 3
241    self.assertEqual(mock_response.mock_calls, expected_calls)
242
243  def testHandleMetadataUpdate(self):
244    mock_response = mock.Mock()
245    mock_response.return_value = {}
246    self.mock_watcher._GetMetadataUpdate = mock_response
247
248    self.assertEqual(self.mock_watcher.GetMetadata(), {})
249    mock_response.assert_called_once_with(
250        metadata_key='', recursive=True, wait=False, timeout=None)
251    self.mock_watcher.logger.exception.assert_not_called()
252
253  def testHandleMetadataUpdateException(self):
254    mock_response = mock.Mock()
255    first = metadata_watcher.socket.timeout()
256    second = metadata_watcher.urlerror.URLError('Test')
257    mock_response.side_effect = [first, first, second, {}]
258    self.mock_watcher._GetMetadataUpdate = mock_response
259    metadata_key = 'instance/id'
260    recursive = False
261    wait = False
262    retry = True
263
264    self.assertEqual(
265        self.mock_watcher._HandleMetadataUpdate(
266            metadata_key=metadata_key, recursive=recursive, wait=wait,
267            timeout=None, retry=retry),
268        {})
269    expected_calls = [
270        mock.call(
271          metadata_key=metadata_key, recursive=recursive, wait=wait,
272          timeout=None),
273    ] * 4
274    self.assertEqual(mock_response.mock_calls, expected_calls)
275    expected_calls = [mock.call.error(mock.ANY, mock.ANY)] * 2
276    self.assertEqual(self.mock_logger.mock_calls, expected_calls)
277
278  def testHandleMetadataUpdateExceptionNoRetry(self):
279    mock_response = mock.Mock()
280    mock_response.side_effect = metadata_watcher.socket.timeout()
281    self.mock_watcher._GetMetadataUpdate = mock_response
282    metadata_key = 'instance/id'
283    recursive = False
284    wait = False
285    retry = False
286
287    self.assertIsNone(
288        self.mock_watcher._HandleMetadataUpdate(
289            metadata_key=metadata_key, recursive=recursive, wait=wait,
290            timeout=None, retry=retry))
291    expected_calls = [
292        mock.call(
293            metadata_key=metadata_key, recursive=recursive, wait=wait,
294            timeout=None),
295    ]
296    self.assertEqual(mock_response.mock_calls, expected_calls)
297    expected_calls = [mock.call.error(mock.ANY, mock.ANY)]
298    self.assertEqual(self.mock_logger.mock_calls, expected_calls)
299
300  def testWatchMetadata(self):
301    mock_response = mock.Mock()
302    mock_response.return_value = {}
303    self.mock_watcher._HandleMetadataUpdate = mock_response
304    mock_handler = mock.Mock()
305    mock_handler.side_effect = Exception()
306    self.mock_logger.exception.side_effect = RuntimeError()
307    recursive = True
308
309    with self.assertRaises(RuntimeError):
310      self.mock_watcher.WatchMetadata(mock_handler, recursive=recursive)
311    mock_handler.assert_called_once_with({})
312    mock_response.assert_called_once_with(
313        metadata_key='', recursive=recursive, wait=True, timeout=None)
314
315  def testWatchMetadataException(self):
316    mock_response = mock.Mock()
317    mock_response.side_effect = metadata_watcher.socket.timeout()
318    self.mock_watcher._GetMetadataUpdate = mock_response
319    self.mock_logger.error.side_effect = RuntimeError()
320    metadata_key = 'instance/id'
321    recursive = False
322
323    with self.assertRaises(RuntimeError):
324      self.mock_watcher.WatchMetadata(
325          None, metadata_key=metadata_key, recursive=recursive)
326    mock_response.assert_called_once_with(
327        metadata_key=metadata_key, recursive=recursive, wait=True, timeout=None)
328
329  def testGetMetadata(self):
330    mock_response = mock.Mock()
331    mock_response.return_value = {}
332    self.mock_watcher._HandleMetadataUpdate = mock_response
333
334    self.assertEqual(self.mock_watcher.GetMetadata(), {})
335    mock_response.assert_called_once_with(
336        metadata_key='', recursive=True, wait=False, timeout=None, retry=True)
337    self.mock_watcher.logger.exception.assert_not_called()
338
339  def testGetMetadataArgs(self):
340    mock_response = mock.Mock()
341    mock_response.return_value = {}
342    self.mock_watcher._HandleMetadataUpdate = mock_response
343    metadata_key = 'instance/id'
344    recursive = False
345    retry = False
346
347    response = self.mock_watcher.GetMetadata(
348        metadata_key=metadata_key, recursive=recursive, timeout=60,
349        retry=retry)
350    self.assertEqual(response, {})
351    mock_response.assert_called_once_with(
352        metadata_key=metadata_key, recursive=False, wait=False, timeout=60,
353        retry=False)
354    self.mock_watcher.logger.exception.assert_not_called()
355
356
357if __name__ == '__main__':
358  unittest.main()
359