functest/utils/decorators.py

   1 #!/usr/bin/env python
   2
   3 # Copyright (c) 2017 Orange and others.
   4 #
   5 # All rights reserved. This program and the accompanying materials
   6 # are made available under the terms of the Apache License, Version 2.0
   7 # which accompanies this distribution, and is available at
   8 # http://www.apache.org/licenses/LICENSE-2.0
   9
  10 # pylint: disable=missing-docstring
  11
  12 import errno
  13 import functools
  14 import os
  15
  16 import mock
  17 import requests.sessions
  18 from six.moves import urllib
  19
  20
  21 def can_dump_request_to_file(method):
  22
  23     def dump_preparedrequest(request, **kwargs):
  24         # pylint: disable=unused-argument
  25         parseresult = urllib.parse.urlparse(request.url)
  26         if parseresult.scheme == "file":
  27             try:
  28                 dirname = os.path.dirname(parseresult.path)
  29                 os.makedirs(dirname)
  30             except OSError as ex:
  31                 if ex.errno != errno.EEXIST:
  32                     raise
  33             with open(parseresult.path, 'a') as dumpfile:
  34                 headers = ""
  35                 for key in request.headers:
  36                     headers += key + " " + request.headers[key] + "\n"
  37                 message = "{} {}\n{}\n{}\n\n\n".format(
  38                     request.method, request.url, headers, request.body)
  39                 dumpfile.write(message)
  40         return mock.Mock()
  41
  42     def patch_request(method, url, **kwargs):
  43         with requests.sessions.Session() as session:
  44             parseresult = urllib.parse.urlparse(url)
  45             if parseresult.scheme == "file":
  46                 with mock.patch.object(session, 'send',
  47                                        side_effect=dump_preparedrequest):
  48                     return session.request(method=method, url=url, **kwargs)
  49             else:
  50                 return session.request(method=method, url=url, **kwargs)
  51
  52     @functools.wraps(method)
  53     def hook(*args, **kwargs):
  54         with mock.patch('requests.api.request', side_effect=patch_request):
  55             return method(*args, **kwargs)
  56
  57     return hook