1
0

helpers.py 57 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659
  1. # -*- coding: utf-8 -*-
  2. '''
  3. :copyright: Copyright 2013-2017 by the SaltStack Team, see AUTHORS for more details.
  4. :license: Apache 2.0, see LICENSE for more details.
  5. tests.support.helpers
  6. ~~~~~~~~~~~~~~~~~~~~~
  7. Test support helpers
  8. '''
  9. # pylint: disable=repr-flag-used-in-string,wrong-import-order
  10. # Import Python libs
  11. from __future__ import absolute_import, print_function, unicode_literals
  12. import base64
  13. import errno
  14. import functools
  15. import inspect
  16. import logging
  17. import os
  18. import random
  19. import shutil
  20. import signal
  21. import socket
  22. import string
  23. import subprocess
  24. import sys
  25. import tempfile
  26. import textwrap
  27. import threading
  28. import time
  29. import tornado.ioloop
  30. import tornado.web
  31. import types
  32. # Import 3rd-party libs
  33. import psutil # pylint: disable=3rd-party-module-not-gated
  34. from salt.ext import six
  35. from salt.ext.six.moves import range, builtins # pylint: disable=import-error,redefined-builtin
  36. try:
  37. from pytestsalt.utils import get_unused_localhost_port # pylint: disable=unused-import
  38. except ImportError:
  39. def get_unused_localhost_port():
  40. '''
  41. Return a random unused port on localhost
  42. '''
  43. usock = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)
  44. usock.bind(('127.0.0.1', 0))
  45. port = usock.getsockname()[1]
  46. usock.close()
  47. return port
  48. # Import Salt Tests Support libs
  49. from tests.support.unit import skip, _id, SkipTest
  50. from tests.support.mock import patch
  51. from tests.support.paths import FILES, TMP
  52. # Import Salt libs
  53. import salt.utils.files
  54. import salt.utils.platform
  55. import salt.utils.stringutils
  56. if salt.utils.platform.is_windows():
  57. import salt.utils.win_functions
  58. else:
  59. import pwd
  60. log = logging.getLogger(__name__)
  61. HAS_SYMLINKS = None
  62. def no_symlinks():
  63. '''
  64. Check if git is installed and has symlinks enabled in the configuration.
  65. '''
  66. global HAS_SYMLINKS
  67. if HAS_SYMLINKS is not None:
  68. return not HAS_SYMLINKS
  69. output = ''
  70. try:
  71. output = subprocess.Popen(
  72. ['git', 'config', '--get', 'core.symlinks'],
  73. cwd=TMP,
  74. stdout=subprocess.PIPE).communicate()[0]
  75. except OSError as exc:
  76. if exc.errno != errno.ENOENT:
  77. raise
  78. except subprocess.CalledProcessError:
  79. # git returned non-zero status
  80. pass
  81. HAS_SYMLINKS = False
  82. if output.strip() == 'true':
  83. HAS_SYMLINKS = True
  84. return not HAS_SYMLINKS
  85. def destructiveTest(caller):
  86. '''
  87. Mark a test case as a destructive test for example adding or removing users
  88. from your system.
  89. .. code-block:: python
  90. class MyTestCase(TestCase):
  91. @destructiveTest
  92. def test_create_user(self):
  93. pass
  94. '''
  95. # Late import
  96. from tests.support.runtests import RUNTIME_VARS
  97. if RUNTIME_VARS.PYTEST_SESSION:
  98. setattr(caller, '__destructive_test__', True)
  99. if inspect.isclass(caller):
  100. # We're decorating a class
  101. old_setup = getattr(caller, 'setUp', None)
  102. def setUp(self, *args, **kwargs):
  103. if os.environ.get('DESTRUCTIVE_TESTS', 'False').lower() == 'false':
  104. self.skipTest('Destructive tests are disabled')
  105. if old_setup is not None:
  106. old_setup(self, *args, **kwargs)
  107. caller.setUp = setUp
  108. return caller
  109. # We're simply decorating functions
  110. @functools.wraps(caller)
  111. def wrap(cls):
  112. if os.environ.get('DESTRUCTIVE_TESTS', 'False').lower() == 'false':
  113. cls.skipTest('Destructive tests are disabled')
  114. return caller(cls)
  115. return wrap
  116. def expensiveTest(caller):
  117. '''
  118. Mark a test case as an expensive test, for example, a test which can cost
  119. money(Salt's cloud provider tests).
  120. .. code-block:: python
  121. class MyTestCase(TestCase):
  122. @expensiveTest
  123. def test_create_user(self):
  124. pass
  125. '''
  126. # Late import
  127. from tests.support.runtests import RUNTIME_VARS
  128. if RUNTIME_VARS.PYTEST_SESSION:
  129. setattr(caller, '__expensive_test__', True)
  130. if inspect.isclass(caller):
  131. # We're decorating a class
  132. old_setup = getattr(caller, 'setUp', None)
  133. def setUp(self, *args, **kwargs):
  134. if os.environ.get('EXPENSIVE_TESTS', 'False').lower() == 'false':
  135. self.skipTest('Expensive tests are disabled')
  136. if old_setup is not None:
  137. old_setup(self, *args, **kwargs)
  138. caller.setUp = setUp
  139. return caller
  140. # We're simply decorating functions
  141. @functools.wraps(caller)
  142. def wrap(cls):
  143. if os.environ.get('EXPENSIVE_TESTS', 'False').lower() == 'false':
  144. cls.skipTest('Expensive tests are disabled')
  145. return caller(cls)
  146. return wrap
  147. def flaky(caller=None, condition=True, attempts=4):
  148. '''
  149. Mark a test as flaky. The test will attempt to run five times,
  150. looking for a successful run. After an immediate second try,
  151. it will use an exponential backoff starting with one second.
  152. .. code-block:: python
  153. class MyTestCase(TestCase):
  154. @flaky
  155. def test_sometimes_works(self):
  156. pass
  157. '''
  158. if caller is None:
  159. return functools.partial(flaky, condition=condition, attempts=attempts)
  160. if isinstance(condition, bool) and condition is False:
  161. # Don't even decorate
  162. return caller
  163. elif callable(condition):
  164. if condition() is False:
  165. # Don't even decorate
  166. return caller
  167. if inspect.isclass(caller):
  168. attrs = [n for n in dir(caller) if n.startswith('test_')]
  169. for attrname in attrs:
  170. try:
  171. function = getattr(caller, attrname)
  172. if not inspect.isfunction(function) and not inspect.ismethod(function):
  173. continue
  174. setattr(caller, attrname, flaky(caller=function, condition=condition, attempts=attempts))
  175. except Exception as exc:
  176. log.exception(exc)
  177. continue
  178. return caller
  179. @functools.wraps(caller)
  180. def wrap(cls):
  181. for attempt in range(0, attempts):
  182. try:
  183. if attempt > 0:
  184. # Run through setUp again
  185. # We only run it after the first iteration(>0) because the regular
  186. # test runner will have already ran setUp the first time
  187. setup = getattr(cls, 'setUp', None)
  188. if callable(setup):
  189. setup()
  190. return caller(cls)
  191. except Exception as exc:
  192. if not isinstance(exc, (AssertionError, SkipTest)) and log.isEnabledFor(logging.DEBUG):
  193. log.exception(exc, exc_info=True)
  194. if attempt >= attempts -1:
  195. # We won't try to run tearDown once the attempts are exhausted
  196. # because the regular test runner will do that for us
  197. six.reraise(*sys.exc_info())
  198. # Run through tearDown again
  199. teardown = getattr(cls, 'tearDown', None)
  200. if callable(teardown):
  201. teardown()
  202. backoff_time = attempt ** 2
  203. log.info(
  204. 'Found Exception. Waiting %s seconds to retry.',
  205. backoff_time
  206. )
  207. time.sleep(backoff_time)
  208. return cls
  209. return wrap
  210. def requires_sshd_server(caller):
  211. '''
  212. Mark a test as requiring the tests SSH daemon running.
  213. .. code-block:: python
  214. class MyTestCase(TestCase):
  215. @requiresSshdServer
  216. def test_create_user(self):
  217. pass
  218. '''
  219. if inspect.isclass(caller):
  220. # We're decorating a class
  221. old_setup = getattr(caller, 'setUp', None)
  222. def setUp(self, *args, **kwargs):
  223. if os.environ.get('SSH_DAEMON_RUNNING', 'False').lower() == 'false':
  224. self.skipTest('SSH tests are disabled')
  225. if old_setup is not None:
  226. old_setup(self, *args, **kwargs)
  227. caller.setUp = setUp
  228. return caller
  229. # We're simply decorating functions
  230. @functools.wraps(caller)
  231. def wrap(cls):
  232. if os.environ.get('SSH_DAEMON_RUNNING', 'False').lower() == 'false':
  233. cls.skipTest('SSH tests are disabled')
  234. return caller(cls)
  235. return wrap
  236. class RedirectStdStreams(object):
  237. '''
  238. Temporarily redirect system output to file like objects.
  239. Default is to redirect to `os.devnull`, which just mutes output, `stdout`
  240. and `stderr`.
  241. '''
  242. def __init__(self, stdout=None, stderr=None):
  243. # Late import
  244. import salt.utils.files
  245. if stdout is None:
  246. stdout = salt.utils.files.fopen(os.devnull, 'w') # pylint: disable=resource-leakage
  247. if stderr is None:
  248. stderr = salt.utils.files.fopen(os.devnull, 'w') # pylint: disable=resource-leakage
  249. self.__stdout = stdout
  250. self.__stderr = stderr
  251. self.__redirected = False
  252. self.patcher = patch.multiple(sys, stderr=self.__stderr, stdout=self.__stdout)
  253. def __enter__(self):
  254. self.redirect()
  255. return self
  256. def __exit__(self, exc_type, exc_value, traceback):
  257. self.unredirect()
  258. def redirect(self):
  259. self.old_stdout = sys.stdout
  260. self.old_stdout.flush()
  261. self.old_stderr = sys.stderr
  262. self.old_stderr.flush()
  263. self.patcher.start()
  264. self.__redirected = True
  265. def unredirect(self):
  266. if not self.__redirected:
  267. return
  268. try:
  269. self.__stdout.flush()
  270. self.__stdout.close()
  271. except ValueError:
  272. # already closed?
  273. pass
  274. try:
  275. self.__stderr.flush()
  276. self.__stderr.close()
  277. except ValueError:
  278. # already closed?
  279. pass
  280. self.patcher.stop()
  281. def flush(self):
  282. if self.__redirected:
  283. try:
  284. self.__stdout.flush()
  285. except Exception:
  286. pass
  287. try:
  288. self.__stderr.flush()
  289. except Exception:
  290. pass
  291. class TstSuiteLoggingHandler(object):
  292. '''
  293. Simple logging handler which can be used to test if certain logging
  294. messages get emitted or not:
  295. .. code-block:: python
  296. with TestsLoggingHandler() as handler:
  297. # (...) Do what ever you wish here
  298. handler.messages # here are the emitted log messages
  299. '''
  300. def __init__(self, level=0, format='%(levelname)s:%(message)s'):
  301. self.level = level
  302. self.format = format
  303. self.activated = False
  304. self.prev_logging_level = None
  305. def activate(self):
  306. class Handler(logging.Handler):
  307. def __init__(self, level):
  308. logging.Handler.__init__(self, level)
  309. self.messages = []
  310. def emit(self, record):
  311. self.messages.append(self.format(record))
  312. self.handler = Handler(self.level)
  313. formatter = logging.Formatter(self.format)
  314. self.handler.setFormatter(formatter)
  315. logging.root.addHandler(self.handler)
  316. self.activated = True
  317. # Make sure we're running with the lowest logging level with our
  318. # tests logging handler
  319. current_logging_level = logging.root.getEffectiveLevel()
  320. if current_logging_level > logging.DEBUG:
  321. self.prev_logging_level = current_logging_level
  322. logging.root.setLevel(0)
  323. def deactivate(self):
  324. if not self.activated:
  325. return
  326. logging.root.removeHandler(self.handler)
  327. # Restore previous logging level if changed
  328. if self.prev_logging_level is not None:
  329. logging.root.setLevel(self.prev_logging_level)
  330. @property
  331. def messages(self):
  332. if not self.activated:
  333. return []
  334. return self.handler.messages
  335. def clear(self):
  336. self.handler.messages = []
  337. def __enter__(self):
  338. self.activate()
  339. return self
  340. def __exit__(self, type, value, traceback):
  341. self.deactivate()
  342. self.activated = False
  343. # Mimic some handler attributes and methods
  344. @property
  345. def lock(self):
  346. if self.activated:
  347. return self.handler.lock
  348. def createLock(self):
  349. if self.activated:
  350. return self.handler.createLock()
  351. def acquire(self):
  352. if self.activated:
  353. return self.handler.acquire()
  354. def release(self):
  355. if self.activated:
  356. return self.handler.release()
  357. def relative_import(import_name, relative_from='../'):
  358. '''
  359. Update sys.path to include `relative_from` before importing `import_name`
  360. '''
  361. try:
  362. return __import__(import_name)
  363. except ImportError:
  364. previous_frame = inspect.getframeinfo(inspect.currentframe().f_back)
  365. sys.path.insert(
  366. 0, os.path.realpath(
  367. os.path.join(
  368. os.path.abspath(
  369. os.path.dirname(previous_frame.filename)
  370. ),
  371. relative_from
  372. )
  373. )
  374. )
  375. return __import__(import_name)
  376. class ForceImportErrorOn(object):
  377. '''
  378. This class is meant to be used in mock'ed test cases which require an
  379. ``ImportError`` to be raised.
  380. >>> import os.path
  381. >>> with ForceImportErrorOn('os.path'):
  382. ... import os.path
  383. ...
  384. Traceback (most recent call last):
  385. File "<stdin>", line 2, in <module>
  386. File "salttesting/helpers.py", line 263, in __import__
  387. 'Forced ImportError raised for \'{0}\''.format(name)
  388. ImportError: Forced ImportError raised for 'os.path'
  389. >>>
  390. >>> with ForceImportErrorOn(('os', 'path')):
  391. ... import os.path
  392. ... sys.modules.pop('os', None)
  393. ... from os import path
  394. ...
  395. <module 'os' from '/usr/lib/python2.7/os.pyc'>
  396. Traceback (most recent call last):
  397. File "<stdin>", line 4, in <module>
  398. File "salttesting/helpers.py", line 288, in __fake_import__
  399. name, ', '.join(fromlist)
  400. ImportError: Forced ImportError raised for 'from os import path'
  401. >>>
  402. >>> with ForceImportErrorOn(('os', 'path'), 'os.path'):
  403. ... import os.path
  404. ... sys.modules.pop('os', None)
  405. ... from os import path
  406. ...
  407. Traceback (most recent call last):
  408. File "<stdin>", line 2, in <module>
  409. File "salttesting/helpers.py", line 281, in __fake_import__
  410. 'Forced ImportError raised for \'{0}\''.format(name)
  411. ImportError: Forced ImportError raised for 'os.path'
  412. >>>
  413. '''
  414. def __init__(self, *module_names):
  415. self.__module_names = {}
  416. for entry in module_names:
  417. if isinstance(entry, (list, tuple)):
  418. modname = entry[0]
  419. self.__module_names[modname] = set(entry[1:])
  420. else:
  421. self.__module_names[entry] = None
  422. self.__original_import = builtins.__import__
  423. self.patcher = patch.object(builtins, '__import__', self.__fake_import__)
  424. def patch_import_function(self):
  425. self.patcher.start()
  426. def restore_import_funtion(self):
  427. self.patcher.stop()
  428. def __fake_import__(self,
  429. name,
  430. globals_={} if six.PY2 else None,
  431. locals_={} if six.PY2 else None,
  432. fromlist=[] if six.PY2 else (),
  433. level=-1 if six.PY2 else 0):
  434. if name in self.__module_names:
  435. importerror_fromlist = self.__module_names.get(name)
  436. if importerror_fromlist is None:
  437. raise ImportError(
  438. 'Forced ImportError raised for \'{0}\''.format(name)
  439. )
  440. if importerror_fromlist.intersection(set(fromlist)):
  441. raise ImportError(
  442. 'Forced ImportError raised for \'{0}\''.format(
  443. 'from {0} import {1}'.format(
  444. name, ', '.join(fromlist)
  445. )
  446. )
  447. )
  448. return self.__original_import(name, globals_, locals_, fromlist, level)
  449. def __enter__(self):
  450. self.patch_import_function()
  451. return self
  452. def __exit__(self, exc_type, exc_value, traceback):
  453. self.restore_import_funtion()
  454. class MockWraps(object):
  455. '''
  456. Helper class to be used with the mock library.
  457. To be used in the ``wraps`` keyword of ``Mock`` or ``MagicMock`` where you
  458. want to trigger a side effect for X times, and afterwards, call the
  459. original and un-mocked method.
  460. As an example:
  461. >>> def original():
  462. ... print 'original'
  463. ...
  464. >>> def side_effect():
  465. ... print 'side effect'
  466. ...
  467. >>> mw = MockWraps(original, 2, side_effect)
  468. >>> mw()
  469. side effect
  470. >>> mw()
  471. side effect
  472. >>> mw()
  473. original
  474. >>>
  475. '''
  476. def __init__(self, original, expected_failures, side_effect):
  477. self.__original = original
  478. self.__expected_failures = expected_failures
  479. self.__side_effect = side_effect
  480. self.__call_counter = 0
  481. def __call__(self, *args, **kwargs):
  482. try:
  483. if self.__call_counter < self.__expected_failures:
  484. if isinstance(self.__side_effect, types.FunctionType):
  485. return self.__side_effect()
  486. raise self.__side_effect
  487. return self.__original(*args, **kwargs)
  488. finally:
  489. self.__call_counter += 1
  490. def requires_network(only_local_network=False):
  491. '''
  492. Simple decorator which is supposed to skip a test case in case there's no
  493. network connection to the internet.
  494. '''
  495. def decorator(func):
  496. @functools.wraps(func)
  497. def wrapper(cls):
  498. has_local_network = False
  499. # First lets try if we have a local network. Inspired in
  500. # verify_socket
  501. try:
  502. pubsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  503. retsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  504. pubsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  505. pubsock.bind(('', 18000))
  506. pubsock.close()
  507. retsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  508. retsock.bind(('', 18001))
  509. retsock.close()
  510. has_local_network = True
  511. except socket.error:
  512. # I wonder if we just have IPV6 support?
  513. try:
  514. pubsock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
  515. retsock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
  516. pubsock.setsockopt(
  517. socket.SOL_SOCKET, socket.SO_REUSEADDR, 1
  518. )
  519. pubsock.bind(('', 18000))
  520. pubsock.close()
  521. retsock.setsockopt(
  522. socket.SOL_SOCKET, socket.SO_REUSEADDR, 1
  523. )
  524. retsock.bind(('', 18001))
  525. retsock.close()
  526. has_local_network = True
  527. except socket.error:
  528. # Let's continue
  529. pass
  530. if only_local_network is True:
  531. if has_local_network is False:
  532. # Since we're only supposed to check local network, and no
  533. # local network was detected, skip the test
  534. cls.skipTest('No local network was detected')
  535. return func(cls)
  536. # We are using the google.com DNS records as numerical IPs to avoid
  537. # DNS lookups which could greatly slow down this check
  538. for addr in ('173.194.41.198', '173.194.41.199', '173.194.41.200',
  539. '173.194.41.201', '173.194.41.206', '173.194.41.192',
  540. '173.194.41.193', '173.194.41.194', '173.194.41.195',
  541. '173.194.41.196', '173.194.41.197'):
  542. sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  543. try:
  544. sock.settimeout(0.25)
  545. sock.connect((addr, 80))
  546. # We connected? Stop the loop
  547. break
  548. except socket.error:
  549. # Let's check the next IP
  550. continue
  551. else:
  552. cls.skipTest('No internet network connection was detected')
  553. finally:
  554. sock.close()
  555. return func(cls)
  556. return wrapper
  557. return decorator
  558. def with_system_user(username, on_existing='delete', delete=True, password=None, groups=None):
  559. '''
  560. Create and optionally destroy a system user to be used within a test
  561. case. The system user is created using the ``user`` salt module.
  562. The decorated testcase function must accept 'username' as an argument.
  563. :param username: The desired username for the system user.
  564. :param on_existing: What to do when the desired username is taken. The
  565. available options are:
  566. * nothing: Do nothing, act as if the user was created.
  567. * delete: delete and re-create the existing user
  568. * skip: skip the test case
  569. '''
  570. if on_existing not in ('nothing', 'delete', 'skip'):
  571. raise RuntimeError(
  572. 'The value of \'on_existing\' can only be one of, '
  573. '\'nothing\', \'delete\' and \'skip\''
  574. )
  575. if not isinstance(delete, bool):
  576. raise RuntimeError(
  577. 'The value of \'delete\' can only be \'True\' or \'False\''
  578. )
  579. def decorator(func):
  580. @functools.wraps(func)
  581. def wrap(cls):
  582. # Let's add the user to the system.
  583. log.debug('Creating system user \'%s\'', username)
  584. kwargs = {'timeout': 60, 'groups': groups}
  585. if salt.utils.platform.is_windows():
  586. kwargs.update({'password': password})
  587. create_user = cls.run_function('user.add', [username], **kwargs)
  588. if not create_user:
  589. log.debug('Failed to create system user')
  590. # The user was not created
  591. if on_existing == 'skip':
  592. cls.skipTest(
  593. 'Failed to create system user \'{0}\''.format(
  594. username
  595. )
  596. )
  597. if on_existing == 'delete':
  598. log.debug('Deleting the system user \'%s\'', username)
  599. delete_user = cls.run_function(
  600. 'user.delete', [username, True, True]
  601. )
  602. if not delete_user:
  603. cls.skipTest(
  604. 'A user named \'{0}\' already existed on the '
  605. 'system and re-creating it was not possible'
  606. .format(username)
  607. )
  608. log.debug('Second time creating system user \'%s\'', username)
  609. create_user = cls.run_function('user.add', [username], **kwargs)
  610. if not create_user:
  611. cls.skipTest(
  612. 'A user named \'{0}\' already existed, was deleted '
  613. 'as requested, but re-creating it was not possible'
  614. .format(username)
  615. )
  616. failure = None
  617. try:
  618. try:
  619. return func(cls, username)
  620. except Exception:
  621. log.exception('Running %s raised an exception', func)
  622. # Store the original exception details which will be raised
  623. # a little further down the code
  624. failure = sys.exc_info()
  625. finally:
  626. if delete:
  627. delete_user = cls.run_function(
  628. 'user.delete', [username, True, True], timeout=60
  629. )
  630. if not delete_user:
  631. if failure is None:
  632. log.warning(
  633. 'Although the actual test-case did not fail, '
  634. 'deleting the created system user \'%s\' '
  635. 'afterwards did.', username
  636. )
  637. else:
  638. log.warning(
  639. 'The test-case failed and also did the removal'
  640. ' of the system user \'%s\'', username
  641. )
  642. if failure is not None:
  643. # If an exception was thrown, raise it
  644. six.reraise(failure[0], failure[1], failure[2])
  645. return wrap
  646. return decorator
  647. def with_system_group(group, on_existing='delete', delete=True):
  648. '''
  649. Create and optionally destroy a system group to be used within a test
  650. case. The system user is crated using the ``group`` salt module.
  651. The decorated testcase function must accept 'group' as an argument.
  652. :param group: The desired group name for the system user.
  653. :param on_existing: What to do when the desired username is taken. The
  654. available options are:
  655. * nothing: Do nothing, act as if the group was created
  656. * delete: delete and re-create the existing user
  657. * skip: skip the test case
  658. '''
  659. if on_existing not in ('nothing', 'delete', 'skip'):
  660. raise RuntimeError(
  661. 'The value of \'on_existing\' can only be one of, '
  662. '\'nothing\', \'delete\' and \'skip\''
  663. )
  664. if not isinstance(delete, bool):
  665. raise RuntimeError(
  666. 'The value of \'delete\' can only be \'True\' or \'False\''
  667. )
  668. def decorator(func):
  669. @functools.wraps(func)
  670. def wrap(cls):
  671. # Let's add the user to the system.
  672. log.debug('Creating system group \'%s\'', group)
  673. create_group = cls.run_function('group.add', [group])
  674. if not create_group:
  675. log.debug('Failed to create system group')
  676. # The group was not created
  677. if on_existing == 'skip':
  678. cls.skipTest(
  679. 'Failed to create system group \'{0}\''.format(group)
  680. )
  681. if on_existing == 'delete':
  682. log.debug('Deleting the system group \'%s\'', group)
  683. delete_group = cls.run_function('group.delete', [group])
  684. if not delete_group:
  685. cls.skipTest(
  686. 'A group named \'{0}\' already existed on the '
  687. 'system and re-creating it was not possible'
  688. .format(group)
  689. )
  690. log.debug('Second time creating system group \'%s\'', group)
  691. create_group = cls.run_function('group.add', [group])
  692. if not create_group:
  693. cls.skipTest(
  694. 'A group named \'{0}\' already existed, was deleted '
  695. 'as requested, but re-creating it was not possible'
  696. .format(group)
  697. )
  698. failure = None
  699. try:
  700. try:
  701. return func(cls, group)
  702. except Exception:
  703. log.exception('Running %s raised an exception', func)
  704. # Store the original exception details which will be raised
  705. # a little further down the code
  706. failure = sys.exc_info()
  707. finally:
  708. if delete:
  709. delete_group = cls.run_function('group.delete', [group])
  710. if not delete_group:
  711. if failure is None:
  712. log.warning(
  713. 'Although the actual test-case did not fail, '
  714. 'deleting the created system group \'%s\' '
  715. 'afterwards did.', group
  716. )
  717. else:
  718. log.warning(
  719. 'The test-case failed and also did the removal'
  720. ' of the system group \'%s\'', group
  721. )
  722. if failure is not None:
  723. # If an exception was thrown, raise it
  724. six.reraise(failure[0], failure[1], failure[2])
  725. return wrap
  726. return decorator
  727. def with_system_user_and_group(username, group,
  728. on_existing='delete', delete=True):
  729. '''
  730. Create and optionally destroy a system user and group to be used within a
  731. test case. The system user is crated using the ``user`` salt module, and
  732. the system group is created with the ``group`` salt module.
  733. The decorated testcase function must accept both the 'username' and 'group'
  734. arguments.
  735. :param username: The desired username for the system user.
  736. :param group: The desired name for the system group.
  737. :param on_existing: What to do when the desired username is taken. The
  738. available options are:
  739. * nothing: Do nothing, act as if the user was created.
  740. * delete: delete and re-create the existing user
  741. * skip: skip the test case
  742. '''
  743. if on_existing not in ('nothing', 'delete', 'skip'):
  744. raise RuntimeError(
  745. 'The value of \'on_existing\' can only be one of, '
  746. '\'nothing\', \'delete\' and \'skip\''
  747. )
  748. if not isinstance(delete, bool):
  749. raise RuntimeError(
  750. 'The value of \'delete\' can only be \'True\' or \'False\''
  751. )
  752. def decorator(func):
  753. @functools.wraps(func)
  754. def wrap(cls):
  755. # Let's add the user to the system.
  756. log.debug('Creating system user %s', username)
  757. create_user = cls.run_function('user.add', [username])
  758. log.debug('Creating system group %s', group)
  759. create_group = cls.run_function('group.add', [group])
  760. if not create_user:
  761. log.debug('Failed to create system user')
  762. # The user was not created
  763. if on_existing == 'skip':
  764. cls.skipTest(
  765. 'Failed to create system user \'{0}\''.format(
  766. username
  767. )
  768. )
  769. if on_existing == 'delete':
  770. log.debug('Deleting the system user \'%s\'', username)
  771. delete_user = cls.run_function(
  772. 'user.delete', [username, True, True]
  773. )
  774. if not delete_user:
  775. cls.skipTest(
  776. 'A user named \'{0}\' already existed on the '
  777. 'system and re-creating it was not possible'
  778. .format(username)
  779. )
  780. log.debug('Second time creating system user \'%s\'', username)
  781. create_user = cls.run_function('user.add', [username])
  782. if not create_user:
  783. cls.skipTest(
  784. 'A user named \'{0}\' already existed, was deleted '
  785. 'as requested, but re-creating it was not possible'
  786. .format(username)
  787. )
  788. if not create_group:
  789. log.debug('Failed to create system group')
  790. # The group was not created
  791. if on_existing == 'skip':
  792. cls.skipTest(
  793. 'Failed to create system group \'{0}\''.format(group)
  794. )
  795. if on_existing == 'delete':
  796. log.debug('Deleting the system group %s', group)
  797. delete_group = cls.run_function('group.delete', [group])
  798. if not delete_group:
  799. cls.skipTest(
  800. 'A group named \'{0}\' already existed on the '
  801. 'system and re-creating it was not possible'
  802. .format(group)
  803. )
  804. log.debug('Second time creating system group %s', group)
  805. create_group = cls.run_function('group.add', [group])
  806. if not create_group:
  807. cls.skipTest(
  808. 'A group named \'{0}\' already existed, was deleted '
  809. 'as requested, but re-creating it was not possible'
  810. .format(group)
  811. )
  812. failure = None
  813. try:
  814. try:
  815. return func(cls, username, group)
  816. except Exception:
  817. log.exception('Running %s raised an exception', func)
  818. # Store the original exception details which will be raised
  819. # a little further down the code
  820. failure = sys.exc_info()
  821. finally:
  822. if delete:
  823. delete_user = cls.run_function(
  824. 'user.delete', [username, True, True]
  825. )
  826. delete_group = cls.run_function('group.delete', [group])
  827. if not delete_user:
  828. if failure is None:
  829. log.warning(
  830. 'Although the actual test-case did not fail, '
  831. 'deleting the created system user \'%s\' '
  832. 'afterwards did.', username
  833. )
  834. else:
  835. log.warning(
  836. 'The test-case failed and also did the removal'
  837. ' of the system user \'%s\'', username
  838. )
  839. if not delete_group:
  840. if failure is None:
  841. log.warning(
  842. 'Although the actual test-case did not fail, '
  843. 'deleting the created system group \'%s\' '
  844. 'afterwards did.', group
  845. )
  846. else:
  847. log.warning(
  848. 'The test-case failed and also did the removal'
  849. ' of the system group \'%s\'', group
  850. )
  851. if failure is not None:
  852. # If an exception was thrown, raise it
  853. six.reraise(failure[0], failure[1], failure[2])
  854. return wrap
  855. return decorator
  856. class WithTempfile(object):
  857. def __init__(self, **kwargs):
  858. self.create = kwargs.pop('create', True)
  859. if 'dir' not in kwargs:
  860. kwargs['dir'] = TMP
  861. if 'prefix' not in kwargs:
  862. kwargs['prefix'] = '__salt.test.'
  863. self.kwargs = kwargs
  864. def __call__(self, func):
  865. self.func = func
  866. return functools.wraps(func)(
  867. lambda testcase, *args, **kwargs: self.wrap(testcase, *args, **kwargs) # pylint: disable=W0108
  868. )
  869. def wrap(self, testcase, *args, **kwargs):
  870. name = salt.utils.files.mkstemp(**self.kwargs)
  871. if not self.create:
  872. os.remove(name)
  873. try:
  874. return self.func(testcase, name, *args, **kwargs)
  875. finally:
  876. try:
  877. os.remove(name)
  878. except OSError:
  879. pass
  880. with_tempfile = WithTempfile
  881. class WithTempdir(object):
  882. def __init__(self, **kwargs):
  883. self.create = kwargs.pop('create', True)
  884. if 'dir' not in kwargs:
  885. kwargs['dir'] = TMP
  886. self.kwargs = kwargs
  887. def __call__(self, func):
  888. self.func = func
  889. return functools.wraps(func)(
  890. lambda testcase, *args, **kwargs: self.wrap(testcase, *args, **kwargs) # pylint: disable=W0108
  891. )
  892. def wrap(self, testcase, *args, **kwargs):
  893. tempdir = tempfile.mkdtemp(**self.kwargs)
  894. if not self.create:
  895. os.rmdir(tempdir)
  896. try:
  897. return self.func(testcase, tempdir, *args, **kwargs)
  898. finally:
  899. shutil.rmtree(tempdir, ignore_errors=True)
  900. with_tempdir = WithTempdir
  901. def requires_system_grains(func):
  902. '''
  903. Function decorator which loads and passes the system's grains to the test
  904. case.
  905. '''
  906. @functools.wraps(func)
  907. def decorator(cls):
  908. if not hasattr(cls, 'run_function'):
  909. raise RuntimeError(
  910. '{0} does not have the \'run_function\' method which is '
  911. 'necessary to collect the system grains'.format(
  912. cls.__class__.__name__
  913. )
  914. )
  915. return func(cls, grains=cls.run_function('grains.items'))
  916. return decorator
  917. def requires_salt_modules(*names):
  918. '''
  919. Makes sure the passed salt module is available. Skips the test if not
  920. .. versionadded:: 0.5.2
  921. '''
  922. def decorator(caller):
  923. if inspect.isclass(caller):
  924. # We're decorating a class
  925. old_setup = getattr(caller, 'setUp', None)
  926. def setUp(self, *args, **kwargs):
  927. if old_setup is not None:
  928. old_setup(self, *args, **kwargs)
  929. if not hasattr(self, 'run_function'):
  930. raise RuntimeError(
  931. '{0} does not have the \'run_function\' method which '
  932. 'is necessary to collect the loaded modules'.format(
  933. self.__class__.__name__
  934. )
  935. )
  936. not_found_modules = self.run_function('runtests_helpers.modules_available', names)
  937. if not_found_modules:
  938. if len(not_found_modules) == 1:
  939. self.skipTest('Salt module \'{0}\' is not available'.format(not_found_modules[0]))
  940. self.skipTest('Salt modules not available: \'{0}\''.format(not_found_modules))
  941. caller.setUp = setUp
  942. return caller
  943. # We're simply decorating functions
  944. @functools.wraps(caller)
  945. def wrapper(cls):
  946. if not hasattr(cls, 'run_function'):
  947. raise RuntimeError(
  948. '{0} does not have the \'run_function\' method which is '
  949. 'necessary to collect the loaded modules'.format(
  950. cls.__class__.__name__
  951. )
  952. )
  953. for name in names:
  954. if name not in cls.run_function('sys.doc', [name]):
  955. cls.skipTest(
  956. 'Salt module \'{0}\' is not available'.format(name)
  957. )
  958. break
  959. return caller(cls)
  960. return wrapper
  961. return decorator
  962. def skip_if_binaries_missing(*binaries, **kwargs):
  963. import salt.utils.path
  964. if len(binaries) == 1:
  965. if isinstance(binaries[0], (list, tuple, set, frozenset)):
  966. binaries = binaries[0]
  967. check_all = kwargs.pop('check_all', False)
  968. message = kwargs.pop('message', None)
  969. if kwargs:
  970. raise RuntimeError(
  971. 'The only supported keyword argument is \'check_all\' and '
  972. '\'message\'. Invalid keyword arguments: {0}'.format(
  973. ', '.join(kwargs.keys())
  974. )
  975. )
  976. if check_all:
  977. for binary in binaries:
  978. if salt.utils.path.which(binary) is None:
  979. return skip(
  980. '{0}The \'{1}\' binary was not found'.format(
  981. message and '{0}. '.format(message) or '',
  982. binary
  983. )
  984. )
  985. elif salt.utils.path.which_bin(binaries) is None:
  986. return skip(
  987. '{0}None of the following binaries was found: {1}'.format(
  988. message and '{0}. '.format(message) or '',
  989. ', '.join(binaries)
  990. )
  991. )
  992. return _id
  993. def skip_if_not_root(func):
  994. # Late import
  995. from tests.support.runtests import RUNTIME_VARS
  996. if RUNTIME_VARS.PYTEST_SESSION:
  997. setattr(func, '__skip_if_not_root__', True)
  998. if not sys.platform.startswith('win'):
  999. if os.getuid() != 0:
  1000. func.__unittest_skip__ = True
  1001. func.__unittest_skip_why__ = 'You must be logged in as root to run this test'
  1002. else:
  1003. current_user = salt.utils.win_functions.get_current_user()
  1004. if current_user != 'SYSTEM':
  1005. if not salt.utils.win_functions.is_admin(current_user):
  1006. func.__unittest_skip__ = True
  1007. func.__unittest_skip_why__ = 'You must be logged in as an Administrator to run this test'
  1008. return func
  1009. if sys.platform.startswith('win'):
  1010. SIGTERM = signal.CTRL_BREAK_EVENT # pylint: disable=no-member
  1011. else:
  1012. SIGTERM = signal.SIGTERM
  1013. def collect_child_processes(pid):
  1014. '''
  1015. Try to collect any started child processes of the provided pid
  1016. '''
  1017. # Let's get the child processes of the started subprocess
  1018. try:
  1019. parent = psutil.Process(pid)
  1020. if hasattr(parent, 'children'):
  1021. children = parent.children(recursive=True)
  1022. else:
  1023. children = []
  1024. except psutil.NoSuchProcess:
  1025. children = []
  1026. return children[::-1] # return a reversed list of the children
  1027. def _terminate_process_list(process_list, kill=False, slow_stop=False):
  1028. for process in process_list[:][::-1]: # Iterate over a reversed copy of the list
  1029. if not psutil.pid_exists(process.pid):
  1030. process_list.remove(process)
  1031. continue
  1032. try:
  1033. if not kill and process.status() == psutil.STATUS_ZOMBIE:
  1034. # Zombie processes will exit once child processes also exit
  1035. continue
  1036. try:
  1037. cmdline = process.cmdline()
  1038. except psutil.AccessDenied:
  1039. # OSX is more restrictive about the above information
  1040. cmdline = None
  1041. except OSError:
  1042. cmdline = None
  1043. if not cmdline:
  1044. try:
  1045. cmdline = process.as_dict()
  1046. except Exception:
  1047. cmdline = 'UNKNOWN PROCESS'
  1048. if kill:
  1049. log.info('Killing process(%s): %s', process.pid, cmdline)
  1050. process.kill()
  1051. else:
  1052. log.info('Terminating process(%s): %s', process.pid, cmdline)
  1053. try:
  1054. if slow_stop:
  1055. # Allow coverage data to be written down to disk
  1056. process.send_signal(SIGTERM)
  1057. try:
  1058. process.wait(2)
  1059. except psutil.TimeoutExpired:
  1060. if psutil.pid_exists(process.pid):
  1061. continue
  1062. else:
  1063. process.terminate()
  1064. except OSError as exc:
  1065. if exc.errno not in (errno.ESRCH, errno.EACCES):
  1066. raise
  1067. if not psutil.pid_exists(process.pid):
  1068. process_list.remove(process)
  1069. except psutil.NoSuchProcess:
  1070. process_list.remove(process)
  1071. def terminate_process_list(process_list, kill=False, slow_stop=False):
  1072. def on_process_terminated(proc):
  1073. log.info('Process %s terminated with exit code: %s', getattr(proc, '_cmdline', proc), proc.returncode)
  1074. # Try to terminate processes with the provided kill and slow_stop parameters
  1075. log.info('Terminating process list. 1st step. kill: %s, slow stop: %s', kill, slow_stop)
  1076. # Cache the cmdline since that will be inaccessible once the process is terminated
  1077. for proc in process_list:
  1078. try:
  1079. cmdline = proc.cmdline()
  1080. except (psutil.NoSuchProcess, psutil.AccessDenied):
  1081. # OSX is more restrictive about the above information
  1082. cmdline = None
  1083. if not cmdline:
  1084. try:
  1085. cmdline = proc
  1086. except (psutil.NoSuchProcess, psutil.AccessDenied):
  1087. cmdline = '<could not be retrived; dead process: {0}>'.format(proc)
  1088. proc._cmdline = cmdline
  1089. _terminate_process_list(process_list, kill=kill, slow_stop=slow_stop)
  1090. psutil.wait_procs(process_list, timeout=15, callback=on_process_terminated)
  1091. if process_list:
  1092. # If there's still processes to be terminated, retry and kill them if slow_stop is False
  1093. log.info('Terminating process list. 2nd step. kill: %s, slow stop: %s', slow_stop is False, slow_stop)
  1094. _terminate_process_list(process_list, kill=slow_stop is False, slow_stop=slow_stop)
  1095. psutil.wait_procs(process_list, timeout=10, callback=on_process_terminated)
  1096. if process_list:
  1097. # If there's still processes to be terminated, just kill them, no slow stopping now
  1098. log.info('Terminating process list. 3rd step. kill: True, slow stop: False')
  1099. _terminate_process_list(process_list, kill=True, slow_stop=False)
  1100. psutil.wait_procs(process_list, timeout=5, callback=on_process_terminated)
  1101. if process_list:
  1102. # In there's still processes to be terminated, log a warning about it
  1103. log.warning('Some processes failed to properly terminate: %s', process_list)
  1104. def terminate_process(pid=None, process=None, children=None, kill_children=False, slow_stop=False):
  1105. '''
  1106. Try to terminate/kill the started processe
  1107. '''
  1108. children = children or []
  1109. process_list = []
  1110. def on_process_terminated(proc):
  1111. if proc.returncode:
  1112. log.info('Process %s terminated with exit code: %s', getattr(proc, '_cmdline', proc), proc.returncode)
  1113. else:
  1114. log.info('Process %s terminated', getattr(proc, '_cmdline', proc))
  1115. if pid and not process:
  1116. try:
  1117. process = psutil.Process(pid)
  1118. process_list.append(process)
  1119. except psutil.NoSuchProcess:
  1120. # Process is already gone
  1121. process = None
  1122. if kill_children:
  1123. if process:
  1124. if not children:
  1125. children = collect_child_processes(process.pid)
  1126. else:
  1127. # Let's collect children again since there might be new ones
  1128. children.extend(collect_child_processes(pid))
  1129. if children:
  1130. process_list.extend(children)
  1131. if process_list:
  1132. if process:
  1133. log.info('Stopping process %s and respective children: %s', process, children)
  1134. else:
  1135. log.info('Terminating process list: %s', process_list)
  1136. terminate_process_list(process_list, kill=slow_stop is False, slow_stop=slow_stop)
  1137. if process and psutil.pid_exists(process.pid):
  1138. log.warning('Process left behind which we were unable to kill: %s', process)
  1139. def terminate_process_pid(pid, only_children=False):
  1140. children = []
  1141. process = None
  1142. # Let's begin the shutdown routines
  1143. try:
  1144. process = psutil.Process(pid)
  1145. children = collect_child_processes(pid)
  1146. except psutil.NoSuchProcess:
  1147. log.info('No process with the PID %s was found running', pid)
  1148. if only_children:
  1149. return terminate_process(children=children, kill_children=True, slow_stop=True)
  1150. return terminate_process(pid=pid, process=process, children=children, kill_children=True, slow_stop=True)
  1151. def repeat(caller=None, condition=True, times=5):
  1152. '''
  1153. Repeat a test X amount of times until the first failure.
  1154. .. code-block:: python
  1155. class MyTestCase(TestCase):
  1156. @repeat
  1157. def test_sometimes_works(self):
  1158. pass
  1159. '''
  1160. if caller is None:
  1161. return functools.partial(repeat, condition=condition, times=times)
  1162. if isinstance(condition, bool) and condition is False:
  1163. # Don't even decorate
  1164. return caller
  1165. elif callable(condition):
  1166. if condition() is False:
  1167. # Don't even decorate
  1168. return caller
  1169. if inspect.isclass(caller):
  1170. attrs = [n for n in dir(caller) if n.startswith('test_')]
  1171. for attrname in attrs:
  1172. try:
  1173. function = getattr(caller, attrname)
  1174. if not inspect.isfunction(function) and not inspect.ismethod(function):
  1175. continue
  1176. setattr(caller, attrname, repeat(caller=function, condition=condition, times=times))
  1177. except Exception as exc:
  1178. log.exception(exc)
  1179. continue
  1180. return caller
  1181. @functools.wraps(caller)
  1182. def wrap(cls):
  1183. result = None
  1184. for attempt in range(1, times+1):
  1185. log.info('%s test run %d of %s times', cls, attempt, times)
  1186. caller(cls)
  1187. return cls
  1188. return wrap
  1189. def http_basic_auth(login_cb=lambda username, password: False):
  1190. '''
  1191. A crude decorator to force a handler to request HTTP Basic Authentication
  1192. Example usage:
  1193. .. code-block:: python
  1194. @http_basic_auth(lambda u, p: u == 'foo' and p == 'bar')
  1195. class AuthenticatedHandler(tornado.web.RequestHandler):
  1196. pass
  1197. '''
  1198. def wrapper(handler_class):
  1199. def wrap_execute(handler_execute):
  1200. def check_auth(handler, kwargs):
  1201. auth = handler.request.headers.get('Authorization')
  1202. if auth is None or not auth.startswith('Basic '):
  1203. # No username/password entered yet, we need to return a 401
  1204. # and set the WWW-Authenticate header to request login.
  1205. handler.set_status(401)
  1206. handler.set_header(
  1207. 'WWW-Authenticate', 'Basic realm=Restricted')
  1208. else:
  1209. # Strip the 'Basic ' from the beginning of the auth header
  1210. # leaving the base64-encoded secret
  1211. username, password = \
  1212. base64.b64decode(auth[6:]).split(':', 1)
  1213. if login_cb(username, password):
  1214. # Authentication successful
  1215. return
  1216. else:
  1217. # Authentication failed
  1218. handler.set_status(403)
  1219. handler._transforms = []
  1220. handler.finish()
  1221. def _execute(self, transforms, *args, **kwargs):
  1222. check_auth(self, kwargs)
  1223. return handler_execute(self, transforms, *args, **kwargs)
  1224. return _execute
  1225. handler_class._execute = wrap_execute(handler_class._execute)
  1226. return handler_class
  1227. return wrapper
  1228. def generate_random_name(prefix, size=6):
  1229. '''
  1230. Generates a random name by combining the provided prefix with a randomly generated
  1231. ascii string.
  1232. .. versionadded:: 2018.3.0
  1233. prefix
  1234. The string to prefix onto the randomly generated ascii string.
  1235. size
  1236. The number of characters to generate. Default: 6.
  1237. '''
  1238. return prefix + ''.join(
  1239. random.choice(string.ascii_uppercase + string.digits)
  1240. for x in range(size)
  1241. )
  1242. class Webserver(object):
  1243. '''
  1244. Starts a tornado webserver on 127.0.0.1 on a random available port
  1245. USAGE:
  1246. .. code-block:: python
  1247. from tests.support.helpers import Webserver
  1248. webserver = Webserver('/path/to/web/root')
  1249. webserver.start()
  1250. webserver.stop()
  1251. '''
  1252. def __init__(self,
  1253. root=None,
  1254. port=None,
  1255. wait=5,
  1256. handler=None):
  1257. '''
  1258. root
  1259. Root directory of webserver. If not passed, it will default to the
  1260. location of the base environment of the integration suite's file
  1261. roots (tests/integration/files/file/base/)
  1262. port
  1263. Port on which to listen. If not passed, a random one will be chosen
  1264. at the time the start() function is invoked.
  1265. wait : 5
  1266. Number of seconds to wait for the socket to be open before raising
  1267. an exception
  1268. handler
  1269. Can be used to use a subclass of tornado.web.StaticFileHandler,
  1270. such as when enforcing authentication with the http_basic_auth
  1271. decorator.
  1272. '''
  1273. if port is not None and not isinstance(port, six.integer_types):
  1274. raise ValueError('port must be an integer')
  1275. if root is None:
  1276. root = os.path.join(FILES, 'file', 'base')
  1277. try:
  1278. self.root = os.path.realpath(root)
  1279. except AttributeError:
  1280. raise ValueError('root must be a string')
  1281. self.port = port
  1282. self.wait = wait
  1283. self.handler = handler \
  1284. if handler is not None \
  1285. else tornado.web.StaticFileHandler
  1286. self.web_root = None
  1287. def target(self):
  1288. '''
  1289. Threading target which stands up the tornado application
  1290. '''
  1291. self.ioloop = tornado.ioloop.IOLoop()
  1292. self.ioloop.make_current()
  1293. self.application = tornado.web.Application(
  1294. [(r'/(.*)', self.handler, {'path': self.root})])
  1295. self.application.listen(self.port)
  1296. self.ioloop.start()
  1297. @property
  1298. def listening(self):
  1299. if self.port is None:
  1300. return False
  1301. sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  1302. return sock.connect_ex(('127.0.0.1', self.port)) == 0
  1303. def url(self, path):
  1304. '''
  1305. Convenience function which, given a file path, will return a URL that
  1306. points to that path. If the path is relative, it will just be appended
  1307. to self.web_root.
  1308. '''
  1309. if self.web_root is None:
  1310. raise RuntimeError('Webserver instance has not been started')
  1311. err_msg = 'invalid path, must be either a relative path or a path ' \
  1312. 'within {0}'.format(self.root)
  1313. try:
  1314. relpath = path \
  1315. if not os.path.isabs(path) \
  1316. else os.path.relpath(path, self.root)
  1317. if relpath.startswith('..' + os.sep):
  1318. raise ValueError(err_msg)
  1319. return '/'.join((self.web_root, relpath))
  1320. except AttributeError:
  1321. raise ValueError(err_msg)
  1322. def start(self):
  1323. '''
  1324. Starts the webserver
  1325. '''
  1326. if self.port is None:
  1327. self.port = get_unused_localhost_port()
  1328. self.web_root = 'http://127.0.0.1:{0}'.format(self.port)
  1329. self.server_thread = threading.Thread(target=self.target)
  1330. self.server_thread.daemon = True
  1331. self.server_thread.start()
  1332. for idx in range(self.wait + 1):
  1333. if self.listening:
  1334. break
  1335. if idx != self.wait:
  1336. time.sleep(1)
  1337. else:
  1338. raise Exception(
  1339. 'Failed to start tornado webserver on 127.0.0.1:{0} within '
  1340. '{1} seconds'.format(self.port, self.wait)
  1341. )
  1342. def stop(self):
  1343. '''
  1344. Stops the webserver
  1345. '''
  1346. self.ioloop.add_callback(self.ioloop.stop)
  1347. self.server_thread.join()
  1348. def win32_kill_process_tree(pid, sig=signal.SIGTERM, include_parent=True,
  1349. timeout=None, on_terminate=None):
  1350. '''
  1351. Kill a process tree (including grandchildren) with signal "sig" and return
  1352. a (gone, still_alive) tuple. "on_terminate", if specified, is a callabck
  1353. function which is called as soon as a child terminates.
  1354. '''
  1355. if pid == os.getpid():
  1356. raise RuntimeError("I refuse to kill myself")
  1357. try:
  1358. parent = psutil.Process(pid)
  1359. except psutil.NoSuchProcess:
  1360. log.debug("PID not found alive: %d", pid)
  1361. return ([], [])
  1362. children = parent.children(recursive=True)
  1363. if include_parent:
  1364. children.append(parent)
  1365. for p in children:
  1366. p.send_signal(sig)
  1367. gone, alive = psutil.wait_procs(children, timeout=timeout,
  1368. callback=on_terminate)
  1369. return (gone, alive)
  1370. def this_user():
  1371. '''
  1372. Get the user associated with the current process.
  1373. '''
  1374. if salt.utils.platform.is_windows():
  1375. return salt.utils.win_functions.get_current_user(with_domain=False)
  1376. return pwd.getpwuid(os.getuid())[0]
  1377. def dedent(text, linesep=os.linesep):
  1378. '''
  1379. A wrapper around textwrap.dedent that also sets line endings.
  1380. '''
  1381. linesep = salt.utils.stringutils.to_unicode(linesep)
  1382. unicode_text = textwrap.dedent(salt.utils.stringutils.to_unicode(text))
  1383. clean_text = linesep.join(unicode_text.splitlines())
  1384. if unicode_text.endswith(u'\n'):
  1385. clean_text += linesep
  1386. if not isinstance(text, six.text_type):
  1387. return salt.utils.stringutils.to_bytes(clean_text)
  1388. return clean_text
  1389. class PatchedEnviron(object):
  1390. def __init__(self, **kwargs):
  1391. self.cleanup_keys = kwargs.pop('__cleanup__', ())
  1392. self.kwargs = kwargs
  1393. self.original_environ = None
  1394. def __enter__(self):
  1395. self.original_environ = os.environ.copy()
  1396. for key in self.cleanup_keys:
  1397. os.environ.pop(key, None)
  1398. # Make sure there are no unicode characters in the self.kwargs if we're
  1399. # on Python 2. These are being added to `os.environ` and causing
  1400. # problems
  1401. if not sys.version_info > (3, 0):
  1402. kwargs = self.kwargs.copy()
  1403. clean_kwargs = {}
  1404. for k in self.kwargs:
  1405. key = k
  1406. if isinstance(key, six.text_type):
  1407. key = key.encode('utf-8')
  1408. if isinstance(self.kwargs[k], six.text_type):
  1409. kwargs[k] = kwargs[k].encode('utf-8')
  1410. clean_kwargs[key] = kwargs[k]
  1411. self.kwargs = clean_kwargs
  1412. os.environ.update(**self.kwargs)
  1413. return self
  1414. def __exit__(self, *args):
  1415. os.environ.clear()
  1416. os.environ.update(self.original_environ)
  1417. patched_environ = PatchedEnviron