__init__.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  1. # -*- coding: utf-8 -*-
  2. # Copyright 2019 New Vector Ltd
  3. # Copyright 2020 The Matrix.org Foundation C.I.C
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. """
  17. Utilities for running the unit tests
  18. """
  19. import sys
  20. import warnings
  21. from asyncio import Future
  22. from typing import Any, Awaitable, Callable, TypeVar
  23. from mock import Mock
  24. import attr
  25. from twisted.python.failure import Failure
  26. from twisted.web.client import ResponseDone
  27. TV = TypeVar("TV")
  28. def get_awaitable_result(awaitable: Awaitable[TV]) -> TV:
  29. """Get the result from an Awaitable which should have completed
  30. Asserts that the given awaitable has a result ready, and returns its value
  31. """
  32. i = awaitable.__await__()
  33. try:
  34. next(i)
  35. except StopIteration as e:
  36. # awaitable returned a result
  37. return e.value
  38. # if next didn't raise, the awaitable hasn't completed.
  39. raise Exception("awaitable has not yet completed")
  40. def make_awaitable(result: Any) -> Awaitable[Any]:
  41. """
  42. Makes an awaitable, suitable for mocking an `async` function.
  43. This uses Futures as they can be awaited multiple times so can be returned
  44. to multiple callers.
  45. """
  46. future = Future() # type: ignore
  47. future.set_result(result)
  48. return future
  49. def setup_awaitable_errors() -> Callable[[], None]:
  50. """
  51. Convert warnings from a non-awaited coroutines into errors.
  52. """
  53. warnings.simplefilter("error", RuntimeWarning)
  54. # unraisablehook was added in Python 3.8.
  55. if not hasattr(sys, "unraisablehook"):
  56. return lambda: None
  57. # State shared between unraisablehook and check_for_unraisable_exceptions.
  58. unraisable_exceptions = []
  59. orig_unraisablehook = sys.unraisablehook # type: ignore
  60. def unraisablehook(unraisable):
  61. unraisable_exceptions.append(unraisable.exc_value)
  62. def cleanup():
  63. """
  64. A method to be used as a clean-up that fails a test-case if there are any new unraisable exceptions.
  65. """
  66. sys.unraisablehook = orig_unraisablehook # type: ignore
  67. if unraisable_exceptions:
  68. raise unraisable_exceptions.pop()
  69. sys.unraisablehook = unraisablehook # type: ignore
  70. return cleanup
  71. def simple_async_mock(return_value=None, raises=None) -> Mock:
  72. # AsyncMock is not available in python3.5, this mimics part of its behaviour
  73. async def cb(*args, **kwargs):
  74. if raises:
  75. raise raises
  76. return return_value
  77. return Mock(side_effect=cb)
  78. @attr.s
  79. class FakeResponse:
  80. """A fake twisted.web.IResponse object
  81. there is a similar class at treq.test.test_response, but it lacks a `phrase`
  82. attribute, and didn't support deliverBody until recently.
  83. """
  84. # HTTP response code
  85. code = attr.ib(type=int)
  86. # HTTP response phrase (eg b'OK' for a 200)
  87. phrase = attr.ib(type=bytes)
  88. # body of the response
  89. body = attr.ib(type=bytes)
  90. def deliverBody(self, protocol):
  91. protocol.dataReceived(self.body)
  92. protocol.connectionLost(Failure(ResponseDone()))