functest/utils/decorators.py

   1 #!/usr/bin/env python
   2
   3 import mock
   4 import requests.sessions
   5 import urlparse
   6
   7
   8 def can_dump_request_to_file(method):
   9
  10     def dump_preparedrequest(request, **kwargs):
  11         parseresult = urlparse.urlparse(request.url)
  12         if parseresult.scheme == "file":
  13             with open(parseresult.path.replace('/results', ''), 'a') as f:
  14                 headers = ""
  15                 for key in request.headers:
  16                     headers += key + " " + request.headers[key] + "\n"
  17                 message = "{} {}\n{}\n{}\n\n\n".format(
  18                     request.method, request.url, headers, request.body)
  19                 f.write(message)
  20         return mock.Mock()
  21
  22     def patch_request(method, url, **kwargs):
  23         with requests.sessions.Session() as session:
  24             parseresult = urlparse.urlparse(url)
  25             if parseresult.scheme == "file":
  26                 with mock.patch.object(
  27                         session, 'send', side_effect=dump_preparedrequest):
  28                     return session.request(method=method, url=url, **kwargs)
  29             else:
  30                 return session.request(method=method, url=url, **kwargs)
  31
  32     def hook(*args, **kwargs):
  33         with mock.patch('requests.api.request', side_effect=patch_request):
  34             return method(*args, **kwargs)
  35
  36     return hook