testing.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870
  1. """reflex.testing - tools for testing reflex apps."""
  2. from __future__ import annotations
  3. import asyncio
  4. import contextlib
  5. import dataclasses
  6. import functools
  7. import inspect
  8. import os
  9. import pathlib
  10. import platform
  11. import re
  12. import signal
  13. import socket
  14. import socketserver
  15. import subprocess
  16. import textwrap
  17. import threading
  18. import time
  19. import types
  20. from http.server import SimpleHTTPRequestHandler
  21. from typing import (
  22. TYPE_CHECKING,
  23. Any,
  24. AsyncIterator,
  25. Callable,
  26. Coroutine,
  27. Optional,
  28. Type,
  29. TypeVar,
  30. Union,
  31. )
  32. import psutil
  33. import uvicorn
  34. import reflex
  35. import reflex.reflex
  36. import reflex.utils.build
  37. import reflex.utils.exec
  38. import reflex.utils.prerequisites
  39. import reflex.utils.processes
  40. from reflex.state import (
  41. BaseState,
  42. StateManagerMemory,
  43. StateManagerRedis,
  44. reload_state_module,
  45. )
  46. try:
  47. from selenium import webdriver # pyright: ignore [reportMissingImports]
  48. from selenium.webdriver.remote.webdriver import ( # pyright: ignore [reportMissingImports]
  49. WebDriver,
  50. )
  51. if TYPE_CHECKING:
  52. from selenium.webdriver.common.options import (
  53. ArgOptions, # pyright: ignore [reportMissingImports]
  54. )
  55. from selenium.webdriver.remote.webelement import ( # pyright: ignore [reportMissingImports]
  56. WebElement,
  57. )
  58. has_selenium = True
  59. except ImportError:
  60. has_selenium = False
  61. # The timeout (minutes) to check for the port.
  62. DEFAULT_TIMEOUT = 15
  63. POLL_INTERVAL = 0.25
  64. FRONTEND_POPEN_ARGS = {}
  65. T = TypeVar("T")
  66. TimeoutType = Optional[Union[int, float]]
  67. if platform.system() == "Windows":
  68. FRONTEND_POPEN_ARGS["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP # type: ignore
  69. FRONTEND_POPEN_ARGS["shell"] = True
  70. else:
  71. FRONTEND_POPEN_ARGS["start_new_session"] = True
  72. # borrowed from py3.11
  73. class chdir(contextlib.AbstractContextManager):
  74. """Non thread-safe context manager to change the current working directory."""
  75. def __init__(self, path):
  76. """Prepare contextmanager.
  77. Args:
  78. path: the path to change to
  79. """
  80. self.path = path
  81. self._old_cwd = []
  82. def __enter__(self):
  83. """Save current directory and perform chdir."""
  84. self._old_cwd.append(os.getcwd())
  85. os.chdir(self.path)
  86. def __exit__(self, *excinfo):
  87. """Change back to previous directory on stack.
  88. Args:
  89. excinfo: sys.exc_info captured in the context block
  90. """
  91. os.chdir(self._old_cwd.pop())
  92. @dataclasses.dataclass
  93. class AppHarness:
  94. """AppHarness executes a reflex app in-process for testing."""
  95. app_name: str
  96. app_source: Optional[
  97. types.FunctionType | types.ModuleType | str | functools.partial
  98. ]
  99. app_path: pathlib.Path
  100. app_module_path: pathlib.Path
  101. app_module: Optional[types.ModuleType] = None
  102. app_instance: Optional[reflex.App] = None
  103. frontend_process: Optional[subprocess.Popen] = None
  104. frontend_url: Optional[str] = None
  105. frontend_output_thread: Optional[threading.Thread] = None
  106. backend_thread: Optional[threading.Thread] = None
  107. backend: Optional[uvicorn.Server] = None
  108. state_manager: Optional[StateManagerMemory | StateManagerRedis] = None
  109. _frontends: list["WebDriver"] = dataclasses.field(default_factory=list)
  110. _decorated_pages: list = dataclasses.field(default_factory=list)
  111. @classmethod
  112. def create(
  113. cls,
  114. root: pathlib.Path,
  115. app_source: Optional[types.FunctionType | types.ModuleType | str] = None,
  116. app_name: Optional[str] = None,
  117. ) -> "AppHarness":
  118. """Create an AppHarness instance at root.
  119. Args:
  120. root: the directory that will contain the app under test.
  121. app_source: if specified, the source code from this function or module is used
  122. as the main module for the app. It may also be the raw source code text, as a str.
  123. If unspecified, then root must already contain a working reflex app and will be used directly.
  124. app_name: provide the name of the app, otherwise will be derived from app_source or root.
  125. Raises:
  126. ValueError: when app_source is a string and app_name is not provided.
  127. Returns:
  128. AppHarness instance
  129. """
  130. if app_name is None:
  131. if app_source is None:
  132. app_name = root.name
  133. elif isinstance(app_source, functools.partial):
  134. keywords = app_source.keywords
  135. slug_suffix = "_".join([str(v) for v in keywords.values()])
  136. func_name = app_source.func.__name__
  137. app_name = f"{func_name}_{slug_suffix}"
  138. app_name = re.sub(r"[^a-zA-Z0-9_]", "_", app_name)
  139. elif isinstance(app_source, str):
  140. raise ValueError(
  141. "app_name must be provided when app_source is a string."
  142. )
  143. else:
  144. app_name = app_source.__name__
  145. app_name = app_name.lower()
  146. while "__" in app_name:
  147. app_name = app_name.replace("__", "_")
  148. return cls(
  149. app_name=app_name,
  150. app_source=app_source,
  151. app_path=root,
  152. app_module_path=root / app_name / f"{app_name}.py",
  153. )
  154. def _get_globals_from_signature(self, func: Any) -> dict[str, Any]:
  155. """Get the globals from a function or module object.
  156. Args:
  157. func: function or module object
  158. Returns:
  159. dict of globals
  160. """
  161. overrides = {}
  162. glbs = {}
  163. if not callable(func):
  164. return glbs
  165. if isinstance(func, functools.partial):
  166. overrides = func.keywords
  167. func = func.func
  168. for param in inspect.signature(func).parameters.values():
  169. if param.default is not inspect.Parameter.empty:
  170. glbs[param.name] = param.default
  171. glbs.update(overrides)
  172. return glbs
  173. def _get_source_from_app_source(self, app_source: Any) -> str:
  174. """Get the source from app_source.
  175. Args:
  176. app_source: function or module or str
  177. Returns:
  178. source code
  179. """
  180. if isinstance(app_source, str):
  181. return app_source
  182. source = inspect.getsource(app_source)
  183. source = re.sub(
  184. r"^\s*def\s+\w+\s*\(.*?\)(\s+->\s+\w+)?:", "", source, flags=re.DOTALL
  185. )
  186. return textwrap.dedent(source)
  187. def _initialize_app(self):
  188. os.environ["TELEMETRY_ENABLED"] = "" # disable telemetry reporting for tests
  189. self.app_path.mkdir(parents=True, exist_ok=True)
  190. if self.app_source is not None:
  191. app_globals = self._get_globals_from_signature(self.app_source)
  192. if isinstance(self.app_source, functools.partial):
  193. self.app_source = self.app_source.func # type: ignore
  194. # get the source from a function or module object
  195. source_code = "\n".join(
  196. [
  197. "\n".join(
  198. self.get_app_global_source(k, v) for k, v in app_globals.items()
  199. ),
  200. self._get_source_from_app_source(self.app_source),
  201. ]
  202. )
  203. with chdir(self.app_path):
  204. reflex.reflex._init(
  205. name=self.app_name,
  206. template=reflex.constants.Templates.DEFAULT,
  207. loglevel=reflex.constants.LogLevel.INFO,
  208. )
  209. self.app_module_path.write_text(source_code)
  210. with chdir(self.app_path):
  211. # ensure config and app are reloaded when testing different app
  212. reflex.config.get_config(reload=True)
  213. # Save decorated pages before importing the test app module
  214. before_decorated_pages = reflex.app.DECORATED_PAGES[self.app_name].copy()
  215. # Ensure the AppHarness test does not skip State assignment due to running via pytest
  216. os.environ.pop(reflex.constants.PYTEST_CURRENT_TEST, None)
  217. self.app_module = reflex.utils.prerequisites.get_compiled_app(reload=True)
  218. # Save the pages that were added during testing
  219. self._decorated_pages = [
  220. p
  221. for p in reflex.app.DECORATED_PAGES[self.app_name]
  222. if p not in before_decorated_pages
  223. ]
  224. self.app_instance = self.app_module.app
  225. if isinstance(self.app_instance._state_manager, StateManagerRedis):
  226. # Create our own redis connection for testing.
  227. self.state_manager = StateManagerRedis.create(self.app_instance.state)
  228. else:
  229. self.state_manager = self.app_instance._state_manager
  230. def _reload_state_module(self):
  231. """Reload the rx.State module to avoid conflict when reloading."""
  232. reload_state_module(module=f"{self.app_name}.{self.app_name}")
  233. def _get_backend_shutdown_handler(self):
  234. if self.backend is None:
  235. raise RuntimeError("Backend was not initialized.")
  236. original_shutdown = self.backend.shutdown
  237. async def _shutdown_redis(*args, **kwargs) -> None:
  238. # ensure redis is closed before event loop
  239. try:
  240. if self.app_instance is not None and isinstance(
  241. self.app_instance.state_manager, StateManagerRedis
  242. ):
  243. await self.app_instance.state_manager.close()
  244. except ValueError:
  245. pass
  246. await original_shutdown(*args, **kwargs)
  247. return _shutdown_redis
  248. def _start_backend(self, port=0):
  249. if self.app_instance is None:
  250. raise RuntimeError("App was not initialized.")
  251. self.backend = uvicorn.Server(
  252. uvicorn.Config(
  253. app=self.app_instance.api,
  254. host="127.0.0.1",
  255. port=port,
  256. )
  257. )
  258. self.backend.shutdown = self._get_backend_shutdown_handler()
  259. self.backend_thread = threading.Thread(target=self.backend.run)
  260. self.backend_thread.start()
  261. def _start_frontend(self):
  262. # Set up the frontend.
  263. with chdir(self.app_path):
  264. config = reflex.config.get_config()
  265. config.api_url = "http://{0}:{1}".format(
  266. *self._poll_for_servers().getsockname(),
  267. )
  268. reflex.utils.build.setup_frontend(self.app_path)
  269. # Start the frontend.
  270. self.frontend_process = reflex.utils.processes.new_process(
  271. [reflex.utils.prerequisites.get_package_manager(), "run", "dev"],
  272. cwd=self.app_path / reflex.constants.Dirs.WEB,
  273. env={"PORT": "0"},
  274. **FRONTEND_POPEN_ARGS,
  275. )
  276. def _wait_frontend(self):
  277. while self.frontend_url is None:
  278. line = (
  279. self.frontend_process.stdout.readline() # pyright: ignore [reportOptionalMemberAccess]
  280. )
  281. if not line:
  282. break
  283. print(line) # for pytest diagnosis
  284. m = re.search(reflex.constants.Next.FRONTEND_LISTENING_REGEX, line)
  285. if m is not None:
  286. self.frontend_url = m.group(1)
  287. break
  288. if self.frontend_url is None:
  289. raise RuntimeError("Frontend did not start")
  290. def consume_frontend_output():
  291. while True:
  292. line = (
  293. self.frontend_process.stdout.readline() # pyright: ignore [reportOptionalMemberAccess]
  294. )
  295. if not line:
  296. break
  297. print(line)
  298. self.frontend_output_thread = threading.Thread(target=consume_frontend_output)
  299. self.frontend_output_thread.start()
  300. def start(self) -> "AppHarness":
  301. """Start the backend in a new thread and dev frontend as a separate process.
  302. Returns:
  303. self
  304. """
  305. self._initialize_app()
  306. self._start_backend()
  307. self._start_frontend()
  308. self._wait_frontend()
  309. return self
  310. @staticmethod
  311. def get_app_global_source(key, value):
  312. """Get the source code of a global object.
  313. If value is a function or class we render the actual
  314. source of value otherwise we assign value to key.
  315. Args:
  316. key: variable name to assign value to.
  317. value: value of the global variable.
  318. Returns:
  319. The rendered app global code.
  320. """
  321. if not inspect.isclass(value) and not inspect.isfunction(value):
  322. return f"{key} = {value!r}"
  323. return inspect.getsource(value)
  324. def __enter__(self) -> "AppHarness":
  325. """Contextmanager protocol for `start()`.
  326. Returns:
  327. Instance of AppHarness after calling start()
  328. """
  329. return self.start()
  330. def stop(self) -> None:
  331. """Stop the frontend and backend servers."""
  332. self._reload_state_module()
  333. if self.backend is not None:
  334. self.backend.should_exit = True
  335. if self.frontend_process is not None:
  336. # https://stackoverflow.com/a/70565806
  337. frontend_children = psutil.Process(self.frontend_process.pid).children(
  338. recursive=True,
  339. )
  340. if platform.system() == "Windows":
  341. self.frontend_process.terminate()
  342. else:
  343. pgrp = os.getpgid(self.frontend_process.pid)
  344. os.killpg(pgrp, signal.SIGTERM)
  345. # kill any remaining child processes
  346. for child in frontend_children:
  347. # It's okay if the process is already gone.
  348. with contextlib.suppress(psutil.NoSuchProcess):
  349. child.terminate()
  350. _, still_alive = psutil.wait_procs(frontend_children, timeout=3)
  351. for child in still_alive:
  352. # It's okay if the process is already gone.
  353. with contextlib.suppress(psutil.NoSuchProcess):
  354. child.kill()
  355. # wait for main process to exit
  356. self.frontend_process.communicate()
  357. if self.backend_thread is not None:
  358. self.backend_thread.join()
  359. if self.frontend_output_thread is not None:
  360. self.frontend_output_thread.join()
  361. for driver in self._frontends:
  362. driver.quit()
  363. # Cleanup decorated pages added during testing
  364. for page in self._decorated_pages:
  365. reflex.app.DECORATED_PAGES[self.app_name].remove(page)
  366. def __exit__(self, *excinfo) -> None:
  367. """Contextmanager protocol for `stop()`.
  368. Args:
  369. excinfo: sys.exc_info captured in the context block
  370. """
  371. self.stop()
  372. @staticmethod
  373. def _poll_for(
  374. target: Callable[[], T],
  375. timeout: TimeoutType = None,
  376. step: TimeoutType = None,
  377. ) -> T | bool:
  378. """Generic polling logic.
  379. Args:
  380. target: callable that returns truthy if polling condition is met.
  381. timeout: max polling time
  382. step: interval between checking target()
  383. Returns:
  384. return value of target() if truthy within timeout
  385. False if timeout elapses
  386. """
  387. if timeout is None:
  388. timeout = DEFAULT_TIMEOUT
  389. if step is None:
  390. step = POLL_INTERVAL
  391. deadline = time.time() + timeout
  392. while time.time() < deadline:
  393. success = target()
  394. if success:
  395. return success
  396. time.sleep(step)
  397. return False
  398. @staticmethod
  399. async def _poll_for_async(
  400. target: Callable[[], Coroutine[None, None, T]],
  401. timeout: TimeoutType = None,
  402. step: TimeoutType = None,
  403. ) -> T | bool:
  404. """Generic polling logic for async functions.
  405. Args:
  406. target: callable that returns truthy if polling condition is met.
  407. timeout: max polling time
  408. step: interval between checking target()
  409. Returns:
  410. return value of target() if truthy within timeout
  411. False if timeout elapses
  412. """
  413. if timeout is None:
  414. timeout = DEFAULT_TIMEOUT
  415. if step is None:
  416. step = POLL_INTERVAL
  417. deadline = time.time() + timeout
  418. while time.time() < deadline:
  419. success = await target()
  420. if success:
  421. return success
  422. await asyncio.sleep(step)
  423. return False
  424. def _poll_for_servers(self, timeout: TimeoutType = None) -> socket.socket:
  425. """Poll backend server for listening sockets.
  426. Args:
  427. timeout: how long to wait for listening socket.
  428. Returns:
  429. first active listening socket on the backend
  430. Raises:
  431. RuntimeError: when the backend hasn't started running
  432. TimeoutError: when server or sockets are not ready
  433. """
  434. if self.backend is None:
  435. raise RuntimeError("Backend is not running.")
  436. backend = self.backend
  437. # check for servers to be initialized
  438. if not self._poll_for(
  439. target=lambda: getattr(backend, "servers", False),
  440. timeout=timeout,
  441. ):
  442. raise TimeoutError("Backend servers are not initialized.")
  443. # check for sockets to be listening
  444. if not self._poll_for(
  445. target=lambda: getattr(backend.servers[0], "sockets", False),
  446. timeout=timeout,
  447. ):
  448. raise TimeoutError("Backend is not listening.")
  449. return backend.servers[0].sockets[0]
  450. def frontend(self, driver_clz: Optional[Type["WebDriver"]] = None) -> "WebDriver":
  451. """Get a selenium webdriver instance pointed at the app.
  452. Args:
  453. driver_clz: webdriver.Chrome (default), webdriver.Firefox, webdriver.Safari,
  454. webdriver.Edge, etc
  455. Returns:
  456. Instance of the given webdriver navigated to the frontend url of the app.
  457. Raises:
  458. RuntimeError: when selenium is not importable or frontend is not running
  459. """
  460. if not has_selenium:
  461. raise RuntimeError(
  462. "Frontend functionality requires `selenium` to be installed, "
  463. "and it could not be imported."
  464. )
  465. if self.frontend_url is None:
  466. raise RuntimeError("Frontend is not running.")
  467. want_headless = False
  468. options: ArgOptions | None = None
  469. if os.environ.get("APP_HARNESS_HEADLESS"):
  470. want_headless = True
  471. if driver_clz is None:
  472. requested_driver = os.environ.get("APP_HARNESS_DRIVER", "Chrome")
  473. driver_clz = getattr(webdriver, requested_driver)
  474. options = getattr(webdriver, f"{requested_driver}Options")()
  475. if driver_clz is webdriver.Chrome and want_headless:
  476. options = webdriver.ChromeOptions()
  477. options.add_argument("--headless=new")
  478. elif driver_clz is webdriver.Firefox and want_headless:
  479. options = webdriver.FirefoxOptions()
  480. options.add_argument("-headless")
  481. elif driver_clz is webdriver.Edge and want_headless:
  482. options = webdriver.EdgeOptions()
  483. options.add_argument("headless")
  484. if options and (args := os.environ.get("APP_HARNESS_DRIVER_ARGS")):
  485. for arg in args.split(","):
  486. options.add_argument(arg)
  487. driver = driver_clz(options=options) # type: ignore
  488. driver.get(self.frontend_url)
  489. self._frontends.append(driver)
  490. return driver
  491. async def get_state(self, token: str) -> BaseState:
  492. """Get the state associated with the given token.
  493. Args:
  494. token: The state token to look up.
  495. Returns:
  496. The state instance associated with the given token
  497. Raises:
  498. RuntimeError: when the app hasn't started running
  499. """
  500. if self.state_manager is None:
  501. raise RuntimeError("state_manager is not set.")
  502. try:
  503. return await self.state_manager.get_state(token)
  504. finally:
  505. if isinstance(self.state_manager, StateManagerRedis):
  506. await self.state_manager.close()
  507. async def set_state(self, token: str, **kwargs) -> None:
  508. """Set the state associated with the given token.
  509. Args:
  510. token: The state token to set.
  511. kwargs: Attributes to set on the state.
  512. Raises:
  513. RuntimeError: when the app hasn't started running
  514. """
  515. if self.state_manager is None:
  516. raise RuntimeError("state_manager is not set.")
  517. state = await self.get_state(token)
  518. for key, value in kwargs.items():
  519. setattr(state, key, value)
  520. try:
  521. await self.state_manager.set_state(token, state)
  522. finally:
  523. if isinstance(self.state_manager, StateManagerRedis):
  524. await self.state_manager.close()
  525. @contextlib.asynccontextmanager
  526. async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
  527. """Modify the state associated with the given token and send update to frontend.
  528. Args:
  529. token: The state token to modify
  530. Yields:
  531. The state instance associated with the given token
  532. Raises:
  533. RuntimeError: when the app hasn't started running
  534. """
  535. if self.state_manager is None:
  536. raise RuntimeError("state_manager is not set.")
  537. if self.app_instance is None:
  538. raise RuntimeError("App is not running.")
  539. app_state_manager = self.app_instance.state_manager
  540. if isinstance(self.state_manager, StateManagerRedis):
  541. # Temporarily replace the app's state manager with our own, since
  542. # the redis connection is on the backend_thread event loop
  543. self.app_instance._state_manager = self.state_manager
  544. try:
  545. async with self.app_instance.modify_state(token) as state:
  546. yield state
  547. finally:
  548. if isinstance(self.state_manager, StateManagerRedis):
  549. self.app_instance._state_manager = app_state_manager
  550. await self.state_manager.close()
  551. def poll_for_content(
  552. self,
  553. element: "WebElement",
  554. timeout: TimeoutType = None,
  555. exp_not_equal: str = "",
  556. ) -> str:
  557. """Poll element.text for change.
  558. Args:
  559. element: selenium webdriver element to check
  560. timeout: how long to poll element.text
  561. exp_not_equal: exit the polling loop when the element text does not match
  562. Returns:
  563. The element text when the polling loop exited
  564. Raises:
  565. TimeoutError: when the timeout expires before text changes
  566. """
  567. if not self._poll_for(
  568. target=lambda: element.text != exp_not_equal,
  569. timeout=timeout,
  570. ):
  571. raise TimeoutError(
  572. f"{element} content remains {exp_not_equal!r} while polling.",
  573. )
  574. return element.text
  575. def poll_for_value(
  576. self,
  577. element: "WebElement",
  578. timeout: TimeoutType = None,
  579. exp_not_equal: str = "",
  580. ) -> Optional[str]:
  581. """Poll element.get_attribute("value") for change.
  582. Args:
  583. element: selenium webdriver element to check
  584. timeout: how long to poll element value attribute
  585. exp_not_equal: exit the polling loop when the value does not match
  586. Returns:
  587. The element value when the polling loop exited
  588. Raises:
  589. TimeoutError: when the timeout expires before value changes
  590. """
  591. if not self._poll_for(
  592. target=lambda: element.get_attribute("value") != exp_not_equal,
  593. timeout=timeout,
  594. ):
  595. raise TimeoutError(
  596. f"{element} content remains {exp_not_equal!r} while polling.",
  597. )
  598. return element.get_attribute("value")
  599. def poll_for_clients(self, timeout: TimeoutType = None) -> dict[str, BaseState]:
  600. """Poll app state_manager for any connected clients.
  601. Args:
  602. timeout: how long to wait for client states
  603. Returns:
  604. active state instances when the polling loop exited
  605. Raises:
  606. RuntimeError: when the app hasn't started running
  607. TimeoutError: when the timeout expires before any states are seen
  608. """
  609. if self.app_instance is None:
  610. raise RuntimeError("App is not running.")
  611. state_manager = self.app_instance.state_manager
  612. assert isinstance(
  613. state_manager, StateManagerMemory
  614. ), "Only works with memory state manager"
  615. if not self._poll_for(
  616. target=lambda: state_manager.states,
  617. timeout=timeout,
  618. ):
  619. raise TimeoutError("No states were observed while polling.")
  620. return state_manager.states
  621. class SimpleHTTPRequestHandlerCustomErrors(SimpleHTTPRequestHandler):
  622. """SimpleHTTPRequestHandler with custom error page handling."""
  623. def __init__(self, *args, error_page_map: dict[int, pathlib.Path], **kwargs):
  624. """Initialize the handler.
  625. Args:
  626. error_page_map: map of error code to error page path
  627. *args: passed through to superclass
  628. **kwargs: passed through to superclass
  629. """
  630. self.error_page_map = error_page_map
  631. super().__init__(*args, **kwargs)
  632. def send_error(
  633. self, code: int, message: str | None = None, explain: str | None = None
  634. ) -> None:
  635. """Send the error page for the given error code.
  636. If the code matches a custom error page, then message and explain are
  637. ignored.
  638. Args:
  639. code: the error code
  640. message: the error message
  641. explain: the error explanation
  642. """
  643. error_page = self.error_page_map.get(code)
  644. if error_page:
  645. self.send_response(code, message)
  646. self.send_header("Connection", "close")
  647. body = error_page.read_bytes()
  648. self.send_header("Content-Type", self.error_content_type)
  649. self.send_header("Content-Length", str(len(body)))
  650. self.end_headers()
  651. self.wfile.write(body)
  652. else:
  653. super().send_error(code, message, explain)
  654. class Subdir404TCPServer(socketserver.TCPServer):
  655. """TCPServer for SimpleHTTPRequestHandlerCustomErrors that serves from a subdir."""
  656. def __init__(
  657. self,
  658. *args,
  659. root: pathlib.Path,
  660. error_page_map: dict[int, pathlib.Path] | None,
  661. **kwargs,
  662. ):
  663. """Initialize the server.
  664. Args:
  665. root: the root directory to serve from
  666. error_page_map: map of error code to error page path
  667. *args: passed through to superclass
  668. **kwargs: passed through to superclass
  669. """
  670. self.root = root
  671. self.error_page_map = error_page_map or {}
  672. super().__init__(*args, **kwargs)
  673. def finish_request(self, request: socket.socket, client_address: tuple[str, int]):
  674. """Finish one request by instantiating RequestHandlerClass.
  675. Args:
  676. request: the requesting socket
  677. client_address: (host, port) referring to the client’s address.
  678. """
  679. self.RequestHandlerClass(
  680. request,
  681. client_address,
  682. self,
  683. directory=str(self.root), # type: ignore
  684. error_page_map=self.error_page_map, # type: ignore
  685. )
  686. class AppHarnessProd(AppHarness):
  687. """AppHarnessProd executes a reflex app in-process for testing.
  688. In prod mode, instead of running `next dev` the app is exported as static
  689. files and served via the builtin python http.server with custom 404 redirect
  690. handling. Additionally, the backend runs in multi-worker mode.
  691. """
  692. frontend_thread: Optional[threading.Thread] = None
  693. frontend_server: Optional[Subdir404TCPServer] = None
  694. def _run_frontend(self):
  695. web_root = self.app_path / reflex.constants.Dirs.WEB_STATIC
  696. error_page_map = {
  697. 404: web_root / "404.html",
  698. }
  699. with Subdir404TCPServer(
  700. ("", 0),
  701. SimpleHTTPRequestHandlerCustomErrors,
  702. root=web_root,
  703. error_page_map=error_page_map,
  704. ) as self.frontend_server:
  705. self.frontend_url = "http://localhost:{1}".format(
  706. *self.frontend_server.socket.getsockname()
  707. )
  708. self.frontend_server.serve_forever()
  709. def _start_frontend(self):
  710. # Set up the frontend.
  711. with chdir(self.app_path):
  712. config = reflex.config.get_config()
  713. config.api_url = "http://{0}:{1}".format(
  714. *self._poll_for_servers().getsockname(),
  715. )
  716. reflex.reflex.export(
  717. zipping=False,
  718. frontend=True,
  719. backend=False,
  720. loglevel=reflex.constants.LogLevel.INFO,
  721. )
  722. self.frontend_thread = threading.Thread(target=self._run_frontend)
  723. self.frontend_thread.start()
  724. def _wait_frontend(self):
  725. self._poll_for(lambda: self.frontend_server is not None)
  726. if self.frontend_server is None or not self.frontend_server.socket.fileno():
  727. raise RuntimeError("Frontend did not start")
  728. def _start_backend(self):
  729. if self.app_instance is None:
  730. raise RuntimeError("App was not initialized.")
  731. os.environ[reflex.constants.SKIP_COMPILE_ENV_VAR] = "yes"
  732. self.backend = uvicorn.Server(
  733. uvicorn.Config(
  734. app=self.app_instance,
  735. host="127.0.0.1",
  736. port=0,
  737. workers=reflex.utils.processes.get_num_workers(),
  738. ),
  739. )
  740. self.backend.shutdown = self._get_backend_shutdown_handler()
  741. self.backend_thread = threading.Thread(target=self.backend.run)
  742. self.backend_thread.start()
  743. def _poll_for_servers(self, timeout: TimeoutType = None) -> socket.socket:
  744. try:
  745. return super()._poll_for_servers(timeout)
  746. finally:
  747. os.environ.pop(reflex.constants.SKIP_COMPILE_ENV_VAR, None)
  748. def stop(self):
  749. """Stop the frontend python webserver."""
  750. super().stop()
  751. if self.frontend_server is not None:
  752. self.frontend_server.shutdown()
  753. if self.frontend_thread is not None:
  754. self.frontend_thread.join()