src/tools/swarming_client/tests/net_utils.py

   1 # Copyright 2014 The Swarming Authors. All rights reserved.
   2 # Use of this source code is governed under the Apache License, Version 2.0 that
   3 # can be found in the LICENSE file.
   4
   5 import logging
   6 import os
   7 import sys
   8 import threading
   9
  10 TEST_DIR = os.path.dirname(os.path.abspath(__file__))
  11 ROOT_DIR = os.path.dirname(TEST_DIR)
  12 sys.path.insert(0, ROOT_DIR)
  13 sys.path.insert(0, os.path.join(ROOT_DIR, 'third_party'))
  14
  15 from depot_tools import auto_stub
  16 from utils import net
  17
  18
  19 class TestCase(auto_stub.TestCase):
  20   """Mocks out url_open() calls."""
  21   def setUp(self):
  22     super(TestCase, self).setUp()
  23     self.mock(net, 'url_open', self._url_open)
  24     self.mock(net, 'url_read_json', self._url_read_json)
  25     self.mock(net, 'sleep_before_retry', lambda *_: None)
  26     self._lock = threading.Lock()
  27     self._requests = []
  28
  29   def tearDown(self):
  30     try:
  31       if not self.has_failed():
  32         self.assertEqual([], self._requests)
  33     finally:
  34       super(TestCase, self).tearDown()
  35
  36   def expected_requests(self, requests):
  37     """Registers the expected requests along their reponses.
  38
  39     Arguments:
  40       request: list of tuple(url, kwargs, response, headers) for normal requests
  41           and tuple(url, kwargs, response) for json requests. kwargs can be a
  42           callable. In that case, it's called with the actual kwargs. It's
  43           useful when the kwargs values are not deterministic.
  44     """
  45     requests = requests[:]
  46     for request in requests:
  47       self.assertEqual(tuple, request.__class__)
  48       # 3 = json request (url_read_json).
  49       # 4 = normal request (url_open).
  50       self.assertIn(len(request), (3, 4))
  51
  52     with self._lock:
  53       self.assertEqual([], self._requests)
  54       self._requests = requests
  55
  56   def _url_open(self, url, **kwargs):
  57     logging.warn('url_open(%s, %s)', url[:500], str(kwargs)[:500])
  58     with self._lock:
  59       if not self._requests:
  60         return None
  61       # Ignore 'stream' argument, it's not important for these tests.
  62       kwargs.pop('stream', None)
  63       for i, n in enumerate(self._requests):
  64         if n[0] == url:
  65           data = self._requests.pop(i)
  66           if len(data) != 4:
  67             self.fail('Expected normal request, got json data; %s' % url)
  68           _, expected_kwargs, result, headers = data
  69           if callable(expected_kwargs):
  70             expected_kwargs(kwargs)
  71           else:
  72             self.assertEqual(expected_kwargs, kwargs)
  73           if result is not None:
  74             return net.HttpResponse.get_fake_response(result, url, headers)
  75           return None
  76     self.fail('Unknown request %s' % url)
  77
  78   def _url_read_json(self, url, **kwargs):
  79     logging.warn('url_read_json(%s, %s)', url[:500], str(kwargs)[:500])
  80     with self._lock:
  81       if not self._requests:
  82         return None
  83       # Ignore 'stream' argument, it's not important for these tests.
  84       kwargs.pop('stream', None)
  85       for i, n in enumerate(self._requests):
  86         if n[0] == url:
  87           data = self._requests.pop(i)
  88           if len(data) != 3:
  89             self.fail('Expected json request, got normal data; %s' % url)
  90           _, expected_kwargs, result = data
  91           if callable(expected_kwargs):
  92             expected_kwargs(kwargs)
  93           else:
  94             self.assertEqual(expected_kwargs, kwargs)
  95           if result is not None:
  96             return result
  97           return None
  98     self.fail('Unknown request %s' % url)