956a526b20
* add support for lifespan tasks * allow passing args to lifespan task * add message to the cancel call * allow asynccontextmanager as lifespan tasks * Fix integration.utils.SessionStorage Previously the SessionStorage util was just looking in localStorage, but the tests didn't catch it because they were asserting the token was not None, rather than asserting it was truthy. Fixed here, because I'm using this structure in the new lifespan test. * If the lifespan task or context takes "app" parameter, pass the FastAPI instance. * test_lifespan: end to end test for register_lifespan_task * In py3.8, Task.cancel takes no args * test_lifespan: use polling to make the test more robust Fix CI failure * Do not allow task_args for better composability --------- Co-authored-by: Masen Furer <m_github@0x26.net>
121 lines
3.7 KiB
Python
121 lines
3.7 KiB
Python
"""Test cases for the FastAPI lifespan integration."""
|
|
from typing import Generator
|
|
|
|
import pytest
|
|
from selenium.webdriver.common.by import By
|
|
|
|
from reflex.testing import AppHarness
|
|
|
|
from .utils import SessionStorage
|
|
|
|
|
|
def LifespanApp():
|
|
"""App with lifespan tasks and context."""
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
|
|
import reflex as rx
|
|
|
|
lifespan_task_global = 0
|
|
lifespan_context_global = 0
|
|
|
|
@asynccontextmanager
|
|
async def lifespan_context(app, inc: int = 1):
|
|
global lifespan_context_global
|
|
print(f"Lifespan context entered: {app}.")
|
|
lifespan_context_global += inc # pyright: ignore[reportUnboundVariable]
|
|
try:
|
|
yield
|
|
finally:
|
|
print("Lifespan context exited.")
|
|
lifespan_context_global += inc
|
|
|
|
async def lifespan_task(inc: int = 1):
|
|
global lifespan_task_global
|
|
print("Lifespan global started.")
|
|
try:
|
|
while True:
|
|
lifespan_task_global += inc # pyright: ignore[reportUnboundVariable]
|
|
await asyncio.sleep(0.1)
|
|
except asyncio.CancelledError as ce:
|
|
print(f"Lifespan global cancelled: {ce}.")
|
|
lifespan_task_global = 0
|
|
|
|
class LifespanState(rx.State):
|
|
@rx.var
|
|
def task_global(self) -> int:
|
|
return lifespan_task_global
|
|
|
|
@rx.var
|
|
def context_global(self) -> int:
|
|
return lifespan_context_global
|
|
|
|
def tick(self, date):
|
|
pass
|
|
|
|
def index():
|
|
return rx.vstack(
|
|
rx.text(LifespanState.task_global, id="task_global"),
|
|
rx.text(LifespanState.context_global, id="context_global"),
|
|
rx.moment(interval=100, on_change=LifespanState.tick),
|
|
)
|
|
|
|
app = rx.App()
|
|
app.register_lifespan_task(lifespan_task)
|
|
app.register_lifespan_task(lifespan_context, inc=2)
|
|
app.add_page(index)
|
|
|
|
|
|
@pytest.fixture()
|
|
def lifespan_app(tmp_path) -> Generator[AppHarness, None, None]:
|
|
"""Start LifespanApp app at tmp_path via AppHarness.
|
|
|
|
Args:
|
|
tmp_path: pytest tmp_path fixture
|
|
|
|
Yields:
|
|
running AppHarness instance
|
|
"""
|
|
with AppHarness.create(
|
|
root=tmp_path,
|
|
app_source=LifespanApp, # type: ignore
|
|
) as harness:
|
|
yield harness
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lifespan(lifespan_app: AppHarness):
|
|
"""Test the lifespan integration.
|
|
|
|
Args:
|
|
lifespan_app: harness for LifespanApp app
|
|
"""
|
|
assert lifespan_app.app_module is not None, "app module is not found"
|
|
assert lifespan_app.app_instance is not None, "app is not running"
|
|
driver = lifespan_app.frontend()
|
|
|
|
ss = SessionStorage(driver)
|
|
assert AppHarness._poll_for(lambda: ss.get("token") is not None), "token not found"
|
|
|
|
context_global = driver.find_element(By.ID, "context_global")
|
|
task_global = driver.find_element(By.ID, "task_global")
|
|
|
|
assert context_global.text == "2"
|
|
assert lifespan_app.app_module.lifespan_context_global == 2 # type: ignore
|
|
|
|
original_task_global_text = task_global.text
|
|
original_task_global_value = int(original_task_global_text)
|
|
lifespan_app.poll_for_content(task_global, exp_not_equal=original_task_global_text)
|
|
assert lifespan_app.app_module.lifespan_task_global > original_task_global_value # type: ignore
|
|
assert int(task_global.text) > original_task_global_value
|
|
|
|
# Kill the backend
|
|
assert lifespan_app.backend is not None
|
|
lifespan_app.backend.should_exit = True
|
|
if lifespan_app.backend_thread is not None:
|
|
lifespan_app.backend_thread.join()
|
|
|
|
# Check that the lifespan tasks have been cancelled
|
|
assert lifespan_app.app_module.lifespan_task_global == 0
|
|
assert lifespan_app.app_module.lifespan_context_global == 4
|