src/tools/swarming_client/tests/net_test.py

   1 #!/usr/bin/env python
   2 # Copyright 2013 The Swarming Authors. All rights reserved.
   3 # Use of this source code is governed under the Apache License, Version 2.0 that
   4 # can be found in the LICENSE file.
   5
   6 # pylint: disable=R0201,W0613
   7
   8 import StringIO
   9 import __builtin__
  10 import contextlib
  11 import logging
  12 import math
  13 import os
  14 import sys
  15 import unittest
  16
  17 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  18 sys.path.insert(0, ROOT_DIR)
  19 sys.path.insert(0, os.path.join(ROOT_DIR, 'third_party'))
  20
  21 from depot_tools import auto_stub
  22 from utils import net
  23
  24
  25 class RetryLoopMockedTest(auto_stub.TestCase):
  26   """Base class for test cases that mock retry loop."""
  27
  28   def setUp(self):
  29     super(RetryLoopMockedTest, self).setUp()
  30     self._retry_attemps_cls = net.RetryAttempt
  31     self.mock(net, 'sleep_before_retry', self.mocked_sleep_before_retry)
  32     self.mock(net, 'current_time', self.mocked_current_time)
  33     self.mock(net, 'RetryAttempt', self.mocked_retry_attempt)
  34     self.sleeps = []
  35     self.attempts = []
  36
  37   def mocked_sleep_before_retry(self, attempt, max_wait):
  38     self.sleeps.append((attempt, max_wait))
  39
  40   def mocked_current_time(self):
  41     # One attempt is one virtual second.
  42     return float(len(self.attempts))
  43
  44   def mocked_retry_attempt(self, *args, **kwargs):
  45     attempt = self._retry_attemps_cls(*args, **kwargs)
  46     self.attempts.append(attempt)
  47     return attempt
  48
  49   def assertAttempts(self, attempts, max_timeout):
  50     """Asserts that retry loop executed given number of |attempts|."""
  51     expected = [(i, max_timeout - i) for i in xrange(attempts)]
  52     actual = [(x.attempt, x.remaining) for x in self.attempts]
  53     self.assertEqual(expected, actual)
  54
  55   def assertSleeps(self, sleeps):
  56     """Asserts that retry loop slept given number of times."""
  57     self.assertEqual(sleeps, len(self.sleeps))
  58
  59
  60 class RetryLoopTest(RetryLoopMockedTest):
  61   """Test for retry_loop implementation."""
  62
  63   def test_sleep_before_retry(self):
  64     # Verifies bounds. Because it's using a pseudo-random number generator and
  65     # not a read random source, it's basically guaranteed to never return the
  66     # same value twice consecutively.
  67     a = net.calculate_sleep_before_retry(0, 0)
  68     b = net.calculate_sleep_before_retry(0, 0)
  69     self.assertTrue(a >= math.pow(1.5, -1), a)
  70     self.assertTrue(b >= math.pow(1.5, -1), b)
  71     self.assertTrue(a < 1.5 + math.pow(1.5, -1), a)
  72     self.assertTrue(b < 1.5 + math.pow(1.5, -1), b)
  73     self.assertNotEqual(a, b)
  74
  75
  76 class GetHttpServiceTest(unittest.TestCase):
  77   """Tests get_http_service implementation."""
  78
  79   def test_get_http_service(self):
  80     def assert_is_appengine_service(service):
  81       """Verifies HttpService is configured for App Engine communications."""
  82       self.assertIsNotNone(service.authenticator)
  83
  84     def assert_is_googlestorage_service(service):
  85       """Verifies HttpService is configured for GS communications."""
  86       self.assertIsNone(service.authenticator)
  87
  88     # Can recognize app engine URLs.
  89     assert_is_appengine_service(
  90         net.get_http_service('https://appengine-app.appspot.com'))
  91     assert_is_appengine_service(
  92         net.get_http_service('https://version-dot-appengine-app.appspot.com'))
  93
  94     # Localhost is also sort of appengine when running on dev server...
  95     assert_is_appengine_service(
  96         net.get_http_service('http://localhost:8080'))
  97
  98     # Check GS urls.
  99     assert_is_googlestorage_service(
 100         net.get_http_service('https://bucket-name.storage.googleapis.com'))
 101
 102
 103 class HttpServiceTest(RetryLoopMockedTest):
 104   """Tests for HttpService class."""
 105
 106   @staticmethod
 107   def mocked_http_service(
 108       url='http://example.com',
 109       perform_request=None,
 110       authorize=None,
 111       login=None):
 112
 113     class MockedAuthenticator(net.Authenticator):
 114       def authorize(self, request):
 115         return authorize(request) if authorize else None
 116       def login(self, allow_user_interaction):
 117         return login(allow_user_interaction) if login else False
 118
 119     class MockedRequestEngine(object):
 120       def perform_request(self, request):
 121         return perform_request(request) if perform_request else None
 122
 123     return net.HttpService(
 124         url,
 125         authenticator=MockedAuthenticator(),
 126         engine=MockedRequestEngine())
 127
 128   def test_request_GET_success(self):
 129     service_url = 'http://example.com'
 130     request_url = '/some_request'
 131     response = 'True'
 132
 133     def mock_perform_request(request):
 134       self.assertTrue(
 135           request.get_full_url().startswith(service_url + request_url))
 136       return request.make_fake_response(response)
 137
 138     service = self.mocked_http_service(url=service_url,
 139         perform_request=mock_perform_request)
 140     self.assertEqual(service.request(request_url).read(), response)
 141     self.assertAttempts(1, net.URL_OPEN_TIMEOUT)
 142
 143   def test_request_POST_success(self):
 144     service_url = 'http://example.com'
 145     request_url = '/some_request'
 146     response = 'True'
 147
 148     def mock_perform_request(request):
 149       self.assertTrue(
 150           request.get_full_url().startswith(service_url + request_url))
 151       self.assertEqual('', request.body)
 152       return request.make_fake_response(response)
 153
 154     service = self.mocked_http_service(url=service_url,
 155         perform_request=mock_perform_request)
 156     self.assertEqual(service.request(request_url, data={}).read(), response)
 157     self.assertAttempts(1, net.URL_OPEN_TIMEOUT)
 158
 159   def test_request_PUT_success(self):
 160     service_url = 'http://example.com'
 161     request_url = '/some_request'
 162     request_body = 'data_body'
 163     response_body = 'True'
 164     content_type = 'application/octet-stream'
 165
 166     def mock_perform_request(request):
 167       self.assertTrue(
 168           request.get_full_url().startswith(service_url + request_url))
 169       self.assertEqual(request_body, request.body)
 170       self.assertEqual(request.method, 'PUT')
 171       self.assertEqual(request.headers['Content-Type'], content_type)
 172       return request.make_fake_response(response_body)
 173
 174     service = self.mocked_http_service(url=service_url,
 175         perform_request=mock_perform_request)
 176     response = service.request(request_url,
 177         data=request_body, content_type=content_type, method='PUT')
 178     self.assertEqual(response.read(), response_body)
 179     self.assertAttempts(1, net.URL_OPEN_TIMEOUT)
 180
 181   def test_request_success_after_failure(self):
 182     response = 'True'
 183     attempts = []
 184
 185     def mock_perform_request(request):
 186       attempts.append(request)
 187       if len(attempts) == 1:
 188         raise net.ConnectionError()
 189       return request.make_fake_response(response)
 190
 191     service = self.mocked_http_service(perform_request=mock_perform_request)
 192     self.assertEqual(service.request('/', data={}).read(), response)
 193     self.assertAttempts(2, net.URL_OPEN_TIMEOUT)
 194
 195   def test_request_failure_max_attempts_default(self):
 196     def mock_perform_request(_request):
 197       raise net.ConnectionError()
 198     service = self.mocked_http_service(perform_request=mock_perform_request)
 199     self.assertEqual(service.request('/'), None)
 200     self.assertAttempts(net.URL_OPEN_MAX_ATTEMPTS, net.URL_OPEN_TIMEOUT)
 201
 202   def test_request_failure_max_attempts(self):
 203     def mock_perform_request(_request):
 204       raise net.ConnectionError()
 205     service = self.mocked_http_service(perform_request=mock_perform_request)
 206     self.assertEqual(service.request('/', max_attempts=23), None)
 207     self.assertAttempts(23, net.URL_OPEN_TIMEOUT)
 208
 209   def test_request_failure_timeout(self):
 210     def mock_perform_request(_request):
 211       raise net.ConnectionError()
 212     service = self.mocked_http_service(perform_request=mock_perform_request)
 213     self.assertEqual(service.request('/', max_attempts=10000), None)
 214     self.assertAttempts(int(net.URL_OPEN_TIMEOUT) + 1, net.URL_OPEN_TIMEOUT)
 215
 216   def test_request_failure_timeout_default(self):
 217     def mock_perform_request(_request):
 218       raise net.ConnectionError()
 219     service = self.mocked_http_service(perform_request=mock_perform_request)
 220     self.assertEqual(service.request('/', timeout=10.), None)
 221     self.assertAttempts(11, 10.0)
 222
 223   def test_request_HTTP_error_no_retry(self):
 224     count = []
 225     def mock_perform_request(request):
 226       count.append(request)
 227       raise net.HttpError(400)
 228
 229     service = self.mocked_http_service(perform_request=mock_perform_request)
 230     self.assertEqual(service.request('/', data={}), None)
 231     self.assertEqual(1, len(count))
 232     self.assertAttempts(1, net.URL_OPEN_TIMEOUT)
 233
 234   def test_request_HTTP_error_retry_404(self):
 235     response = 'data'
 236     attempts = []
 237
 238     def mock_perform_request(request):
 239       attempts.append(request)
 240       if len(attempts) == 1:
 241         raise net.HttpError(404)
 242       return request.make_fake_response(response)
 243
 244     service = self.mocked_http_service(perform_request=mock_perform_request)
 245     result = service.request('/', data={}, retry_404=True)
 246     self.assertEqual(result.read(), response)
 247     self.assertAttempts(2, net.URL_OPEN_TIMEOUT)
 248
 249   def test_request_HTTP_error_with_retry(self):
 250     response = 'response'
 251     attempts = []
 252
 253     def mock_perform_request(request):
 254       attempts.append(request)
 255       if len(attempts) == 1:
 256         raise net.HttpError(500)
 257       return request.make_fake_response(response)
 258
 259     service = self.mocked_http_service(perform_request=mock_perform_request)
 260     self.assertTrue(service.request('/', data={}).read(), response)
 261     self.assertAttempts(2, net.URL_OPEN_TIMEOUT)
 262
 263   def test_auth_success(self):
 264     calls = []
 265     response = 'response'
 266
 267     def mock_perform_request(request):
 268       calls.append('request')
 269       if 'login' not in calls:
 270         raise net.HttpError(403)
 271       return request.make_fake_response(response)
 272
 273     def mock_authorize(request):
 274       calls.append('authorize')
 275
 276     def mock_login(allow_user_interaction):
 277       self.assertFalse(allow_user_interaction)
 278       calls.append('login')
 279       return True
 280
 281     service = self.mocked_http_service(
 282         perform_request=mock_perform_request,
 283         authorize=mock_authorize,
 284         login=mock_login)
 285     self.assertEqual(service.request('/').read(), response)
 286     self.assertEqual(
 287         ['authorize', 'request', 'login', 'authorize', 'request'], calls)
 288     self.assertAttempts(2, net.URL_OPEN_TIMEOUT)
 289     self.assertSleeps(0)
 290
 291   def test_auth_failure(self):
 292     count = []
 293
 294     def mock_perform_request(_request):
 295       raise net.HttpError(403)
 296
 297     def mock_login(allow_user_interaction):
 298       self.assertFalse(allow_user_interaction)
 299       count.append(1)
 300       return False
 301
 302     service = self.mocked_http_service(perform_request=mock_perform_request,
 303         login=mock_login)
 304     self.assertEqual(service.request('/'), None)
 305     self.assertEqual(len(count), 1)
 306     self.assertAttempts(1, net.URL_OPEN_TIMEOUT)
 307
 308   def test_url_read(self):
 309     # Successfully reads the data.
 310     self.mock(net, 'url_open',
 311         lambda url, **_kwargs: net.HttpResponse.get_fake_response('111', url))
 312     self.assertEqual(net.url_read('https://fake_url.com/test'), '111')
 313
 314     # Respects url_open connection errors.
 315     self.mock(net, 'url_open', lambda _url, **_kwargs: None)
 316     self.assertIsNone(net.url_read('https://fake_url.com/test'))
 317
 318     # Respects read timeout errors.
 319     def timeouting_http_response(url):
 320       def read_mock(_size=None):
 321         raise net.TimeoutError()
 322       response = net.HttpResponse.get_fake_response('', url)
 323       self.mock(response, 'read', read_mock)
 324       return response
 325
 326     self.mock(net, 'url_open',
 327         lambda url, **_kwargs: timeouting_http_response(url))
 328     self.assertIsNone(net.url_read('https://fake_url.com/test'))
 329
 330   def test_url_retrieve(self):
 331     # Successfully reads the data.
 332     @contextlib.contextmanager
 333     def fake_open(_filepath, _mode):
 334       yield StringIO.StringIO()
 335
 336     self.mock(__builtin__, 'open', fake_open)
 337     self.mock(net, 'url_open',
 338         lambda url, **_kwargs: net.HttpResponse.get_fake_response('111', url))
 339     self.assertEqual(
 340         True, net.url_retrieve('filepath', 'https://localhost/test'))
 341
 342     # Respects url_open connection errors.
 343     self.mock(net, 'url_open', lambda _url, **_kwargs: None)
 344     self.assertEqual(
 345         False, net.url_retrieve('filepath', 'https://localhost/test'))
 346
 347     # Respects read timeout errors.
 348     def timeouting_http_response(url):
 349       def read_mock(_size=None):
 350         raise net.TimeoutError()
 351       response = net.HttpResponse.get_fake_response('', url)
 352       self.mock(response, 'read', read_mock)
 353       return response
 354
 355     removed = []
 356     self.mock(os, 'remove', removed.append)
 357     self.mock(net, 'url_open',
 358         lambda url, **_kwargs: timeouting_http_response(url))
 359     self.assertEqual(
 360         False, net.url_retrieve('filepath', 'https://localhost/test'))
 361     self.assertEqual(['filepath'], removed)
 362
 363
 364 if __name__ == '__main__':
 365   logging.basicConfig(
 366       level=(logging.DEBUG if '-v' in sys.argv else logging.FATAL))
 367   if '-v' in sys.argv:
 368     unittest.TestCase.maxDiff = None
 369   unittest.main()