__init__.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204
  1. # -*- encoding: utf8 -*-
  2. #
  3. # The Qubes OS Project, http://www.qubes-os.org
  4. #
  5. # Copyright (C) 2017 Marek Marczykowski-Górecki
  6. # <marmarek@invisiblethingslab.com>
  7. #
  8. # This program is free software; you can redistribute it and/or modify
  9. # it under the terms of the GNU Lesser General Public License as published by
  10. # the Free Software Foundation; either version 2.1 of the License, or
  11. # (at your option) any later version.
  12. #
  13. # This program is distributed in the hope that it will be useful,
  14. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  16. # GNU Lesser General Public License for more details.
  17. #
  18. # You should have received a copy of the GNU Lesser General Public License along
  19. # with this program; if not, see <http://www.gnu.org/licenses/>.
  20. import subprocess
  21. import traceback
  22. import unittest
  23. import io
  24. import qubesadmin
  25. import qubesadmin.app
  26. class TestVM(object):
  27. def __init__(self, name, **kwargs):
  28. self.name = name
  29. self.klass = 'TestVM'
  30. for key, value in kwargs.items():
  31. setattr(self, key, value)
  32. def get_power_state(self):
  33. return getattr(self, 'power_state', 'Running')
  34. def __str__(self):
  35. return self.name
  36. def __lt__(self, other):
  37. if isinstance(other, TestVM):
  38. return self.name < other.name
  39. return NotImplemented
  40. class TestVMCollection(dict):
  41. def __iter__(self):
  42. return iter(self.values())
  43. class TestProcess(object):
  44. def __init__(self, input_callback=None, stdout=None, stderr=None):
  45. self.input_callback = input_callback
  46. self.stdin = io.BytesIO()
  47. # don't let anyone close it, before we get the value
  48. self.stdin_close = self.stdin.close
  49. if self.input_callback:
  50. self.stdin.close = (
  51. lambda: self.input_callback(self.stdin.getvalue()))
  52. else:
  53. self.stdin.close = lambda: None
  54. if stdout == subprocess.PIPE:
  55. self.stdout = io.BytesIO()
  56. else:
  57. self.stdout = stdout
  58. if stderr == subprocess.PIPE:
  59. self.stderr = io.BytesIO()
  60. else:
  61. self.stderr = stderr
  62. self.returncode = 0
  63. def communicate(self, input=None):
  64. if input is not None:
  65. self.stdin.write(input)
  66. self.stdin.close()
  67. self.stdin_close()
  68. return self.stdout, self.stderr
  69. def wait(self):
  70. self.stdin_close()
  71. return 0
  72. def poll(self):
  73. return None
  74. class _AssertNotRaisesContext(object):
  75. """A context manager used to implement TestCase.assertNotRaises methods.
  76. Stolen from unittest and hacked. Regexp support stripped.
  77. """ # pylint: disable=too-few-public-methods
  78. def __init__(self, expected, test_case, expected_regexp=None):
  79. if expected_regexp is not None:
  80. raise NotImplementedError('expected_regexp is unsupported')
  81. self.expected = expected
  82. self.exception = None
  83. self.failureException = test_case.failureException
  84. def __enter__(self):
  85. return self
  86. def __exit__(self, exc_type, exc_value, tb):
  87. if exc_type is None:
  88. return True
  89. if issubclass(exc_type, self.expected):
  90. raise self.failureException(
  91. "{!r} raised, traceback:\n{!s}".format(
  92. exc_value, ''.join(traceback.format_tb(tb))))
  93. else:
  94. # pass through
  95. return False
  96. class QubesTest(qubesadmin.app.QubesBase):
  97. expected_calls = None
  98. actual_calls = None
  99. service_calls = None
  100. def __init__(self):
  101. super(QubesTest, self).__init__()
  102. #: expected calls and saved replies for them
  103. self.expected_calls = {}
  104. #: actual calls made
  105. self.actual_calls = []
  106. #: rpc service calls
  107. self.service_calls = []
  108. def qubesd_call(self, dest, method, arg=None, payload=None,
  109. payload_stream=None):
  110. if payload_stream:
  111. payload = (payload or b'') + payload_stream.read()
  112. call_key = (dest, method, arg, payload)
  113. self.actual_calls.append(call_key)
  114. if call_key not in self.expected_calls:
  115. raise AssertionError('Unexpected call {!r}'.format(call_key))
  116. return_data = self.expected_calls[call_key]
  117. if isinstance(return_data, list):
  118. try:
  119. return_data = return_data.pop(0)
  120. except IndexError:
  121. raise AssertionError('Extra call {!r}'.format(call_key))
  122. return self._parse_qubesd_response(return_data)
  123. def run_service(self, dest, service, **kwargs):
  124. self.service_calls.append((dest, service, kwargs))
  125. return TestProcess(lambda input: self.service_calls.append((dest,
  126. service, input)),
  127. stdout=kwargs.get('stdout', None),
  128. stderr=kwargs.get('stderr', None),
  129. )
  130. class QubesTestCase(unittest.TestCase):
  131. def setUp(self):
  132. super(QubesTestCase, self).setUp()
  133. self.app = QubesTest()
  134. def assertAllCalled(self):
  135. self.assertEqual(
  136. set(self.app.expected_calls.keys()),
  137. set(self.app.actual_calls))
  138. # and also check if calls expected multiple times were called
  139. self.assertFalse([(call, ret)
  140. for call, ret in self.app.expected_calls.items() if
  141. isinstance(ret, list) and ret],
  142. 'Some calls not called expected number of times')
  143. def assertNotRaises(self, excClass, callableObj=None, *args, **kwargs):
  144. """Fail if an exception of class excClass is raised
  145. by callableObj when invoked with arguments args and keyword
  146. arguments kwargs. If a different type of exception is
  147. raised, it will not be caught, and the test case will be
  148. deemed to have suffered an error, exactly as for an
  149. unexpected exception.
  150. If called with callableObj omitted or None, will return a
  151. context object used like this::
  152. with self.assertRaises(SomeException):
  153. do_something()
  154. The context manager keeps a reference to the exception as
  155. the 'exception' attribute. This allows you to inspect the
  156. exception after the assertion::
  157. with self.assertRaises(SomeException) as cm:
  158. do_something()
  159. the_exception = cm.exception
  160. self.assertEqual(the_exception.error_code, 3)
  161. """
  162. context = _AssertNotRaisesContext(excClass, self)
  163. if callableObj is None:
  164. return context
  165. with context:
  166. callableObj(*args, **kwargs)