selenium.py 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263
  1. import sys
  2. import unittest
  3. from contextlib import contextmanager
  4. from functools import wraps
  5. from pathlib import Path
  6. from django.conf import settings
  7. from django.test import LiveServerTestCase, override_settings, tag
  8. from django.utils.functional import classproperty
  9. from django.utils.module_loading import import_string
  10. from django.utils.text import capfirst
  11. class SeleniumTestCaseBase(type(LiveServerTestCase)):
  12. # List of browsers to dynamically create test classes for.
  13. browsers = []
  14. # A selenium hub URL to test against.
  15. selenium_hub = None
  16. # The external host Selenium Hub can reach.
  17. external_host = None
  18. # Sentinel value to differentiate browser-specific instances.
  19. browser = None
  20. # Run browsers in headless mode.
  21. headless = False
  22. def __new__(cls, name, bases, attrs):
  23. """
  24. Dynamically create new classes and add them to the test module when
  25. multiple browsers specs are provided (e.g. --selenium=firefox,chrome).
  26. """
  27. test_class = super().__new__(cls, name, bases, attrs)
  28. # If the test class is either browser-specific or a test base, return it.
  29. if test_class.browser or not any(
  30. name.startswith("test") and callable(value) for name, value in attrs.items()
  31. ):
  32. return test_class
  33. elif test_class.browsers:
  34. # Reuse the created test class to make it browser-specific.
  35. # We can't rename it to include the browser name or create a
  36. # subclass like we do with the remaining browsers as it would
  37. # either duplicate tests or prevent pickling of its instances.
  38. first_browser = test_class.browsers[0]
  39. test_class.browser = first_browser
  40. # Listen on an external interface if using a selenium hub.
  41. host = test_class.host if not test_class.selenium_hub else "0.0.0.0"
  42. test_class.host = host
  43. test_class.external_host = cls.external_host
  44. # Create subclasses for each of the remaining browsers and expose
  45. # them through the test's module namespace.
  46. module = sys.modules[test_class.__module__]
  47. for browser in test_class.browsers[1:]:
  48. browser_test_class = cls.__new__(
  49. cls,
  50. "%s%s" % (capfirst(browser), name),
  51. (test_class,),
  52. {
  53. "browser": browser,
  54. "host": host,
  55. "external_host": cls.external_host,
  56. "__module__": test_class.__module__,
  57. },
  58. )
  59. setattr(module, browser_test_class.__name__, browser_test_class)
  60. return test_class
  61. # If no browsers were specified, skip this class (it'll still be discovered).
  62. return unittest.skip("No browsers specified.")(test_class)
  63. @classmethod
  64. def import_webdriver(cls, browser):
  65. return import_string("selenium.webdriver.%s.webdriver.WebDriver" % browser)
  66. @classmethod
  67. def import_options(cls, browser):
  68. return import_string("selenium.webdriver.%s.options.Options" % browser)
  69. @classmethod
  70. def get_capability(cls, browser):
  71. from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
  72. return getattr(DesiredCapabilities, browser.upper())
  73. def create_options(self):
  74. options = self.import_options(self.browser)()
  75. if self.headless:
  76. match self.browser:
  77. case "chrome" | "edge":
  78. options.add_argument("--headless=new")
  79. case "firefox":
  80. options.add_argument("-headless")
  81. return options
  82. def create_webdriver(self):
  83. options = self.create_options()
  84. if self.selenium_hub:
  85. from selenium import webdriver
  86. for key, value in self.get_capability(self.browser).items():
  87. options.set_capability(key, value)
  88. return webdriver.Remote(command_executor=self.selenium_hub, options=options)
  89. return self.import_webdriver(self.browser)(options=options)
  90. class ChangeWindowSize:
  91. def __init__(self, width, height, selenium):
  92. self.selenium = selenium
  93. self.new_size = (width, height)
  94. def __enter__(self):
  95. self.old_size = self.selenium.get_window_size()
  96. self.selenium.set_window_size(*self.new_size)
  97. return self
  98. def __exit__(self, exc_type, exc_value, traceback):
  99. self.selenium.set_window_size(self.old_size["width"], self.old_size["height"])
  100. @tag("selenium")
  101. class SeleniumTestCase(LiveServerTestCase, metaclass=SeleniumTestCaseBase):
  102. implicit_wait = 10
  103. external_host = None
  104. screenshots = False
  105. @classmethod
  106. def __init_subclass__(cls, **kwargs):
  107. super().__init_subclass__(**kwargs)
  108. if not cls.screenshots:
  109. return
  110. for name, func in list(cls.__dict__.items()):
  111. if not hasattr(func, "_screenshot_cases"):
  112. continue
  113. # Remove the main test.
  114. delattr(cls, name)
  115. # Add separate tests for each screenshot type.
  116. for screenshot_case in getattr(func, "_screenshot_cases"):
  117. @wraps(func)
  118. def test(self, *args, _func=func, _case=screenshot_case, **kwargs):
  119. with getattr(self, _case)():
  120. return _func(self, *args, **kwargs)
  121. test.__name__ = f"{name}_{screenshot_case}"
  122. test.__qualname__ = f"{test.__qualname__}_{screenshot_case}"
  123. test._screenshot_name = name
  124. test._screenshot_case = screenshot_case
  125. setattr(cls, test.__name__, test)
  126. @classproperty
  127. def live_server_url(cls):
  128. return "http://%s:%s" % (cls.external_host or cls.host, cls.server_thread.port)
  129. @classproperty
  130. def allowed_host(cls):
  131. return cls.external_host or cls.host
  132. @classmethod
  133. def setUpClass(cls):
  134. cls.selenium = cls.create_webdriver()
  135. cls.selenium.implicitly_wait(cls.implicit_wait)
  136. super().setUpClass()
  137. cls.addClassCleanup(cls._quit_selenium)
  138. @contextmanager
  139. def desktop_size(self):
  140. with ChangeWindowSize(1280, 720, self.selenium):
  141. yield
  142. @contextmanager
  143. def small_screen_size(self):
  144. with ChangeWindowSize(1024, 768, self.selenium):
  145. yield
  146. @contextmanager
  147. def mobile_size(self):
  148. with ChangeWindowSize(360, 800, self.selenium):
  149. yield
  150. @contextmanager
  151. def rtl(self):
  152. with self.desktop_size():
  153. with override_settings(LANGUAGE_CODE=settings.LANGUAGES_BIDI[-1]):
  154. yield
  155. @contextmanager
  156. def dark(self):
  157. # Navigate to a page before executing a script.
  158. self.selenium.get(self.live_server_url)
  159. self.selenium.execute_script("localStorage.setItem('theme', 'dark');")
  160. with self.desktop_size():
  161. try:
  162. yield
  163. finally:
  164. self.selenium.execute_script("localStorage.removeItem('theme');")
  165. def set_emulated_media(self, *, media=None, features=None):
  166. if self.browser not in {"chrome", "edge"}:
  167. self.skipTest(
  168. "Emulation.setEmulatedMedia is only supported on Chromium and "
  169. "Chrome-based browsers. See https://chromedevtools.github.io/devtools-"
  170. "protocol/1-3/Emulation/#method-setEmulatedMedia for more details."
  171. )
  172. params = {}
  173. if media is not None:
  174. params["media"] = media
  175. if features is not None:
  176. params["features"] = features
  177. # Not using .execute_cdp_cmd() as it isn't supported by the remote web driver
  178. # when using --selenium-hub.
  179. self.selenium.execute(
  180. driver_command="executeCdpCommand",
  181. params={"cmd": "Emulation.setEmulatedMedia", "params": params},
  182. )
  183. @contextmanager
  184. def high_contrast(self):
  185. self.set_emulated_media(features=[{"name": "forced-colors", "value": "active"}])
  186. with self.desktop_size():
  187. try:
  188. yield
  189. finally:
  190. self.set_emulated_media(
  191. features=[{"name": "forced-colors", "value": "none"}]
  192. )
  193. def take_screenshot(self, name):
  194. if not self.screenshots:
  195. return
  196. test = getattr(self, self._testMethodName)
  197. filename = f"{test._screenshot_name}--{name}--{test._screenshot_case}.png"
  198. path = Path.cwd() / "screenshots" / filename
  199. path.parent.mkdir(exist_ok=True, parents=True)
  200. self.selenium.save_screenshot(path)
  201. @classmethod
  202. def _quit_selenium(cls):
  203. # quit() the WebDriver before attempting to terminate and join the
  204. # single-threaded LiveServerThread to avoid a dead lock if the browser
  205. # kept a connection alive.
  206. if hasattr(cls, "selenium"):
  207. cls.selenium.quit()
  208. @contextmanager
  209. def disable_implicit_wait(self):
  210. """Disable the default implicit wait."""
  211. self.selenium.implicitly_wait(0)
  212. try:
  213. yield
  214. finally:
  215. self.selenium.implicitly_wait(self.implicit_wait)
  216. def screenshot_cases(method_names):
  217. if isinstance(method_names, str):
  218. method_names = method_names.split(",")
  219. def wrapper(func):
  220. func._screenshot_cases = method_names
  221. setattr(func, "tags", {"screenshot"}.union(getattr(func, "tags", set())))
  222. return func
  223. return wrapper