client/tests/net_test.py

#!/usr/bin/env vpython3
# Copyright 2013 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

import contextlib
import io
import math
import os
import unittest

# Mutates sys.path.
import test_env

from parameterized import parameterized

# third_party/
from depot_tools import auto_stub

import net_utils

from utils import authenticators
from utils import net


class RetryLoopMockedTest(auto_stub.TestCase):
  """Base class for test cases that mock retry loop."""

  def setUp(self):
    super(RetryLoopMockedTest, self).setUp()
    self._retry_attemps_cls = net.RetryAttempt
    self.mock(net, 'sleep_before_retry', self.mocked_sleep_before_retry)
    self.mock(net, 'current_time', self.mocked_current_time)
    self.mock(net, 'RetryAttempt', self.mocked_retry_attempt)
    self.sleeps = []
    self.attempts = []

  def mocked_sleep_before_retry(self, attempt, max_wait):
    self.sleeps.append((attempt, max_wait))

  def mocked_current_time(self):
    # One attempt is one virtual second.
    return float(len(self.attempts))

  def mocked_retry_attempt(self, *args, **kwargs):
    attempt = self._retry_attemps_cls(*args, **kwargs)
    self.attempts.append(attempt)
    return attempt

  def assertAttempts(self, attempts, max_timeout):
    """Asserts that retry loop executed given number of |attempts|."""
    expected = [(i, max_timeout - i) for i in range(attempts)]
    actual = [(x.attempt, x.remaining) for x in self.attempts]
    self.assertEqual(expected, actual)

  def assertSleeps(self, sleeps):
    """Asserts that retry loop slept given number of times."""
    self.assertEqual(sleeps, len(self.sleeps))


class RetryLoopTest(RetryLoopMockedTest):
  """Test for retry_loop implementation."""

  def test_sleep_before_retry(self):
    # Verifies bounds. Because it's using a pseudo-random number generator and
    # not a read random source, it's basically guaranteed to never return the
    # same value twice consecutively.
    a = net.calculate_sleep_before_retry(0, 0)
    b = net.calculate_sleep_before_retry(0, 0)
    self.assertTrue(a >= math.pow(1.5, -1), a)
    self.assertTrue(b >= math.pow(1.5, -1), b)
    self.assertTrue(a < 1.5 + math.pow(1.5, -1), a)
    self.assertTrue(b < 1.5 + math.pow(1.5, -1), b)
    self.assertNotEqual(a, b)


class HttpErrorTest(unittest.TestCase):
  def test_description_simple(self):
    e = net_utils.make_fake_error(
        code=404,
        url='/url',
        content=b'Response\nBody',
        headers={
            'Header-A': 'value-a',
            'X-Skipped': 'zzz'
        })
    self.assertEqual(e.description(False), 'Server returned HTTP code 404')
    self.assertEqual(e.description(True), '\n'.join([
        'Server returned HTTP code 404',
        '----------',
        'Content-length: 13',
        'Header-a: value-a',
        '',
        'Response',
        'Body',
        '----------',
    ]))

  def test_description_json_good(self):
    e = net_utils.make_fake_error(
        code=404,
        url='/url',
        content=b'{"error": "inner error message"}',
        headers={'Content-Type': 'application/json'})
    self.assertEqual(
        e.description(False),
        'Server returned HTTP code 404 - inner error message')
    self.assertEqual(e.description(True), '\n'.join([
        'Server returned HTTP code 404 - inner error message',
        '----------',
        'Content-length: 32',
        'Content-type: application/json',
        '',
        '{"error": "inner error message"}',
        '----------',
    ]))


class HttpServiceTest(RetryLoopMockedTest):
  """Tests for HttpService class."""

  @staticmethod
  def mocked_http_service(
      url='http://example.com',
      perform_request=None,
      authorize=None,
      login=None):

    class MockedAuthenticator(authenticators.Authenticator):
      supports_login = True
      def authorize(self, request):
        return authorize(request) if authorize else None
      def login(self, allow_user_interaction):
        return login(allow_user_interaction) if login else False

    class MockedRequestEngine:

      @net.maybe_inject_user_agent
      def perform_request(self, request):
        return perform_request(request) if perform_request else None

    return net.HttpService(
        url,
        authenticator=MockedAuthenticator(),
        engine=MockedRequestEngine())

  def test_request_GET_success(self):
    service_url = 'http://example.com'
    request_url = '/some_request'
    response = b'True'

    def mock_perform_request(request):
      self.assertTrue(
          request.get_full_url().startswith(service_url + request_url))
      return net_utils.make_fake_response(response, request.get_full_url())

    service = self.mocked_http_service(url=service_url,
        perform_request=mock_perform_request)
    self.assertEqual(service.request(request_url).read(), response)
    self.assertAttempts(1, net.URL_OPEN_TIMEOUT)

  def test_request_POST_success(self):
    service_url = 'http://example.com'
    request_url = '/some_request'
    response = b'True'

    def mock_perform_request(request):
      self.assertTrue(
          request.get_full_url().startswith(service_url + request_url))
      self.assertEqual(b'', request.body)
      return net_utils.make_fake_response(response, request.get_full_url())

    service = self.mocked_http_service(url=service_url,
        perform_request=mock_perform_request)
    self.assertEqual(service.request(request_url, data={}).read(), response)
    self.assertAttempts(1, net.URL_OPEN_TIMEOUT)

  def test_request_PUT_success(self):
    service_url = 'http://example.com'
    request_url = '/some_request'
    request_body = b'data_body'
    response_body = b'True'
    content_type = 'application/octet-stream'

    def mock_perform_request(request):
      self.assertTrue(
          request.get_full_url().startswith(service_url + request_url))
      self.assertEqual(request_body, request.body)
      self.assertEqual(request.method, 'PUT')
      self.assertEqual(request.headers['Content-Type'], content_type)
      return net_utils.make_fake_response(response_body, request.get_full_url())

    service = self.mocked_http_service(url=service_url,
        perform_request=mock_perform_request)
    response = service.request(request_url,
        data=request_body, content_type=content_type, method='PUT')
    self.assertEqual(response.read(), response_body)
    self.assertAttempts(1, net.URL_OPEN_TIMEOUT)

  def test_request_success_after_failure(self):
    response = b'True'
    attempts = []

    def mock_perform_request(request):
      attempts.append(request)
      if len(attempts) == 1:
        raise net.ConnectionError()
      return net_utils.make_fake_response(response, request.get_full_url())

    service = self.mocked_http_service(perform_request=mock_perform_request)
    self.assertEqual(service.request('/', data={}).read(), response)
    self.assertAttempts(2, net.URL_OPEN_TIMEOUT)

  def test_request_failure_max_attempts_default(self):
    def mock_perform_request(_request):
      raise net.ConnectionError()
    service = self.mocked_http_service(perform_request=mock_perform_request)
    self.assertEqual(service.request('/'), None)
    self.assertAttempts(net.URL_OPEN_MAX_ATTEMPTS, net.URL_OPEN_TIMEOUT)

  def test_request_failure_max_attempts(self):
    def mock_perform_request(_request):
      raise net.ConnectionError()
    service = self.mocked_http_service(perform_request=mock_perform_request)
    self.assertEqual(service.request('/', max_attempts=23), None)
    self.assertAttempts(23, net.URL_OPEN_TIMEOUT)

  def test_request_failure_timeout(self):
    def mock_perform_request(_request):
      raise net.ConnectionError()
    service = self.mocked_http_service(perform_request=mock_perform_request)
    self.assertEqual(service.request('/', max_attempts=10000), None)
    self.assertAttempts(int(net.URL_OPEN_TIMEOUT) + 1, net.URL_OPEN_TIMEOUT)

  def test_request_failure_timeout_default(self):
    def mock_perform_request(_request):
      raise net.ConnectionError()
    service = self.mocked_http_service(perform_request=mock_perform_request)
    self.assertEqual(service.request('/', timeout=10.), None)
    self.assertAttempts(11, 10.0)

  def test_request_HTTP_error_no_retry(self):
    count = []
    def mock_perform_request(request):
      count.append(request)
      raise net_utils.make_fake_error(400, request.get_full_url())

    service = self.mocked_http_service(perform_request=mock_perform_request)
    self.assertEqual(service.request('/', data={}), None)
    self.assertEqual(1, len(count))
    self.assertAttempts(1, net.URL_OPEN_TIMEOUT)

  def test_request_expected_HTTP_error_code(self):
    content = b'data'

    count = []
    def mock_perform_request(request):
      count.append(request)
      raise net_utils.make_fake_error(
          400, request.get_full_url(), content=content)

    service = self.mocked_http_service(perform_request=mock_perform_request)
    response = service.request('/', data={}, expected_error_codes=(400,))
    self.assertEqual(400, response.code)
    self.assertEqual(content, response.read())
    self.assertEqual(1, len(count))
    self.assertAttempts(1, net.URL_OPEN_TIMEOUT)

  def test_request_HTTP_error_retry_404_endpoints(self):
    response = b'data'
    attempts = []

    def mock_perform_request(request):
      attempts.append(request)
      if len(attempts) == 1:
        raise net_utils.make_fake_error(404, request.get_full_url(), headers={
          'Content-Type': 'application/text; charset=ASCII',
        })
      return net_utils.make_fake_response(response, request.get_full_url())

    service = self.mocked_http_service(perform_request=mock_perform_request)
    result = service.request('/_ah/api/foo/v1/bar', )
    self.assertEqual(result.read(), response)
    self.assertAttempts(2, net.URL_OPEN_TIMEOUT)

  @parameterized.expand([(408, ), (429, ), (500, ), (503, ), (599, )])
  def test_request_HTTP_error_with_retry(self, status_code):
    response = b'response'
    attempts = []

    def mock_perform_request(request):
      attempts.append(request)
      if len(attempts) == 1:
        raise net_utils.make_fake_error(status_code, request.get_full_url())
      return net_utils.make_fake_response(response, request.get_full_url())

    service = self.mocked_http_service(perform_request=mock_perform_request)
    self.assertTrue(service.request('/', data={}).read(), response)
    self.assertAttempts(2, net.URL_OPEN_TIMEOUT)

  def test_auth_success(self):
    calls = []
    response = b'response'

    def mock_perform_request(request):
      calls.append('request')
      if 'login' not in calls:
        raise net_utils.make_fake_error(403, request.get_full_url())
      return net_utils.make_fake_response(response, request.get_full_url())

    def mock_authorize(request):
      self.assertTrue(request)
      calls.append('authorize')

    def mock_login(allow_user_interaction):
      self.assertFalse(allow_user_interaction)
      calls.append('login')
      return True

    service = self.mocked_http_service(
        perform_request=mock_perform_request,
        authorize=mock_authorize,
        login=mock_login)
    self.assertEqual(service.request('/').read(), response)
    self.assertEqual(
        ['authorize', 'request', 'login', 'authorize', 'request'], calls)
    self.assertAttempts(2, net.URL_OPEN_TIMEOUT)
    self.assertSleeps(0)

  def test_auth_failure(self):
    count = []

    def mock_perform_request(request):
      raise net_utils.make_fake_error(403, request.get_full_url())

    def mock_login(allow_user_interaction):
      self.assertFalse(allow_user_interaction)
      count.append(1)
      return False

    service = self.mocked_http_service(perform_request=mock_perform_request,
        login=mock_login)
    self.assertEqual(service.request('/'), None)
    self.assertEqual(len(count), 1)
    self.assertAttempts(1, net.URL_OPEN_TIMEOUT)

  def test_url_read(self):
    # Successfully reads the data.
    self.mock(net, 'url_open',
              lambda url, **_kwargs: net_utils.make_fake_response(b'111', url))
    self.assertEqual(net.url_read('https://fake_url.com/test'), b'111')

    # Respects url_open connection errors.
    self.mock(net, 'url_open', lambda _url, **_kwargs: None)
    self.assertIsNone(net.url_read('https://fake_url.com/test'))

    # Respects read timeout errors.
    def timeouting_http_response(url):
      def read_mock(_size=None):
        raise net.TimeoutError()

      response = net_utils.make_fake_response(b'', url)
      self.mock(response, 'read', read_mock)
      return response

    self.mock(net, 'url_open',
        lambda url, **_kwargs: timeouting_http_response(url))
    self.assertIsNone(net.url_read('https://fake_url.com/test'))

  def test_url_retrieve(self):
    # Successfully reads the data.
    @contextlib.contextmanager
    def fake_open(_filepath, _mode):
      yield io.BytesIO()

    self.mock(io, 'open', fake_open)
    self.mock(net, 'url_open',
              lambda url, **_kwargs: net_utils.make_fake_response(b'111', url))
    self.assertEqual(
        True, net.url_retrieve('filepath', 'https://localhost/test'))

    # Respects url_open connection errors.
    self.mock(net, 'url_open', lambda _url, **_kwargs: None)
    self.assertEqual(
        False, net.url_retrieve('filepath', 'https://localhost/test'))

    # Respects read timeout errors.
    def timeouting_http_response(url):
      def iter_content_mock(_size=None):
        raise net.TimeoutError()

      response = net_utils.make_fake_response(b'', url)
      self.mock(response, 'iter_content', iter_content_mock)
      return response

    removed = []
    self.mock(os, 'remove', removed.append)
    self.mock(net, 'url_open',
        lambda url, **_kwargs: timeouting_http_response(url))
    self.assertEqual(
        False, net.url_retrieve('filepath', 'https://localhost/test'))
    self.assertEqual(['filepath'], removed)

  def test_request_has_user_agent(self):

    def mock_perform_request(request):
      return request

    service = self.mocked_http_service(perform_request=mock_perform_request)
    user_agent = 'Py/Foo'
    net.set_user_agent(user_agent)
    request_headers = service.request('/', data={}).headers
    self.assertEqual(request_headers['User-Agent'], user_agent)

    net.set_user_agent(None)
    request_headers = service.request('/', data={}).headers
    self.assertNotIn('User-Agent', request_headers)


class TestNetFunctions(auto_stub.TestCase):
  def test_fix_url(self):
    data = [
      ('http://foo.com/', 'http://foo.com'),
      ('https://foo.com/', 'https://foo.com'),
      ('https://foo.com', 'https://foo.com'),
      ('https://foo.com/a', 'https://foo.com/a'),
      ('https://foo.com/a/', 'https://foo.com/a'),
      ('https://foo.com:8080/a/', 'https://foo.com:8080/a'),
      ('foo.com', 'https://foo.com'),
      ('foo.com/', 'https://foo.com'),
      ('foo.com/a/', 'https://foo.com/a'),
    ]
    for value, expected in data:
      self.assertEqual(expected, net.fix_url(value))


if __name__ == '__main__':
  test_env.main()