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