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