test_state.py 61 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111
  1. from __future__ import annotations
  2. import asyncio
  3. import copy
  4. import datetime
  5. import functools
  6. import json
  7. import os
  8. import sys
  9. from typing import Dict, Generator, List
  10. from unittest.mock import AsyncMock, Mock
  11. import pytest
  12. from plotly.graph_objects import Figure
  13. import reflex as rx
  14. from reflex.base import Base
  15. from reflex.constants import APP_VAR, IS_HYDRATED, RouteVar, SocketEvent
  16. from reflex.event import Event, EventHandler
  17. from reflex.state import (
  18. ImmutableStateError,
  19. LockExpiredError,
  20. MutableProxy,
  21. State,
  22. StateManager,
  23. StateManagerMemory,
  24. StateManagerRedis,
  25. StateProxy,
  26. StateUpdate,
  27. )
  28. from reflex.utils import format, prerequisites
  29. from reflex.vars import BaseVar, ComputedVar
  30. from .states import GenState
  31. CI = bool(os.environ.get("CI", False))
  32. LOCK_EXPIRATION = 2000 if CI else 100
  33. LOCK_EXPIRE_SLEEP = 2.5 if CI else 0.2
  34. class Object(Base):
  35. """A test object fixture."""
  36. prop1: int = 42
  37. prop2: str = "hello"
  38. class TestState(State):
  39. """A test state."""
  40. # Set this class as not test one
  41. __test__ = False
  42. num1: int
  43. num2: float = 3.14
  44. key: str
  45. map_key: str = "a"
  46. array: List[float] = [1, 2, 3.14]
  47. mapping: Dict[str, List[int]] = {"a": [1, 2, 3], "b": [4, 5, 6]}
  48. obj: Object = Object()
  49. complex: Dict[int, Object] = {1: Object(), 2: Object()}
  50. fig: Figure = Figure()
  51. dt: datetime.datetime = datetime.datetime.fromisoformat("1989-11-09T18:53:00+01:00")
  52. @ComputedVar
  53. def sum(self) -> float:
  54. """Dynamically sum the numbers.
  55. Returns:
  56. The sum of the numbers.
  57. """
  58. return self.num1 + self.num2
  59. @ComputedVar
  60. def upper(self) -> str:
  61. """Uppercase the key.
  62. Returns:
  63. The uppercased key.
  64. """
  65. return self.key.upper()
  66. def do_something(self):
  67. """Do something."""
  68. pass
  69. class ChildState(TestState):
  70. """A child state fixture."""
  71. value: str
  72. count: int = 23
  73. def change_both(self, value: str, count: int):
  74. """Change both the value and count.
  75. Args:
  76. value: The new value.
  77. count: The new count.
  78. """
  79. self.value = value.upper()
  80. self.count = count * 2
  81. class ChildState2(TestState):
  82. """A child state fixture."""
  83. value: str
  84. class GrandchildState(ChildState):
  85. """A grandchild state fixture."""
  86. value2: str
  87. def do_nothing(self):
  88. """Do something."""
  89. pass
  90. @pytest.fixture
  91. def test_state() -> TestState:
  92. """A state.
  93. Returns:
  94. A test state.
  95. """
  96. return TestState() # type: ignore
  97. @pytest.fixture
  98. def child_state(test_state) -> ChildState:
  99. """A child state.
  100. Args:
  101. test_state: A test state.
  102. Returns:
  103. A test child state.
  104. """
  105. child_state = test_state.get_substate(["child_state"])
  106. assert child_state is not None
  107. return child_state
  108. @pytest.fixture
  109. def child_state2(test_state) -> ChildState2:
  110. """A second child state.
  111. Args:
  112. test_state: A test state.
  113. Returns:
  114. A second test child state.
  115. """
  116. child_state2 = test_state.get_substate(["child_state2"])
  117. assert child_state2 is not None
  118. return child_state2
  119. @pytest.fixture
  120. def grandchild_state(child_state) -> GrandchildState:
  121. """A state.
  122. Args:
  123. child_state: A child state.
  124. Returns:
  125. A test state.
  126. """
  127. grandchild_state = child_state.get_substate(["grandchild_state"])
  128. assert grandchild_state is not None
  129. return grandchild_state
  130. def test_base_class_vars(test_state):
  131. """Test that the class vars are set correctly.
  132. Args:
  133. test_state: A state.
  134. """
  135. fields = test_state.get_fields()
  136. cls = type(test_state)
  137. for field in fields:
  138. if field in test_state.get_skip_vars():
  139. continue
  140. prop = getattr(cls, field)
  141. assert isinstance(prop, BaseVar)
  142. assert prop.name == field
  143. assert cls.num1.type_ == int
  144. assert cls.num2.type_ == float
  145. assert cls.key.type_ == str
  146. def test_computed_class_var(test_state):
  147. """Test that the class computed vars are set correctly.
  148. Args:
  149. test_state: A state.
  150. """
  151. cls = type(test_state)
  152. vars = [(prop.name, prop.type_) for prop in cls.computed_vars.values()]
  153. assert ("sum", float) in vars
  154. assert ("upper", str) in vars
  155. def test_class_vars(test_state):
  156. """Test that the class vars are set correctly.
  157. Args:
  158. test_state: A state.
  159. """
  160. cls = type(test_state)
  161. assert set(cls.vars.keys()) == {
  162. IS_HYDRATED, # added by hydrate_middleware to all State
  163. "num1",
  164. "num2",
  165. "key",
  166. "map_key",
  167. "array",
  168. "mapping",
  169. "obj",
  170. "complex",
  171. "sum",
  172. "upper",
  173. "fig",
  174. "dt",
  175. }
  176. def test_event_handlers(test_state):
  177. """Test that event handler is set correctly.
  178. Args:
  179. test_state: A state.
  180. """
  181. expected = {
  182. "do_something",
  183. "set_array",
  184. "set_complex",
  185. "set_fig",
  186. "set_key",
  187. "set_mapping",
  188. "set_num1",
  189. "set_num2",
  190. "set_obj",
  191. }
  192. cls = type(test_state)
  193. assert set(cls.event_handlers.keys()).intersection(expected) == expected
  194. def test_default_value(test_state):
  195. """Test that the default value of a var is correct.
  196. Args:
  197. test_state: A state.
  198. """
  199. assert test_state.num1 == 0
  200. assert test_state.num2 == 3.14
  201. assert test_state.key == ""
  202. assert test_state.sum == 3.14
  203. assert test_state.upper == ""
  204. def test_computed_vars(test_state):
  205. """Test that the computed var is computed correctly.
  206. Args:
  207. test_state: A state.
  208. """
  209. test_state.num1 = 1
  210. test_state.num2 = 4
  211. assert test_state.sum == 5
  212. test_state.key = "hello world"
  213. assert test_state.upper == "HELLO WORLD"
  214. def test_dict(test_state):
  215. """Test that the dict representation of a state is correct.
  216. Args:
  217. test_state: A state.
  218. """
  219. substates = {"child_state", "child_state2"}
  220. assert set(test_state.dict().keys()) == set(test_state.vars.keys()) | substates
  221. assert (
  222. set(test_state.dict(include_computed=False).keys())
  223. == set(test_state.base_vars) | substates
  224. )
  225. def test_format_state(test_state):
  226. """Test that the format state is correct.
  227. Args:
  228. test_state: A state.
  229. """
  230. formatted_state = format.format_state(test_state.dict())
  231. exp_formatted_state = {
  232. "array": [1, 2, 3.14],
  233. "child_state": {"count": 23, "grandchild_state": {"value2": ""}, "value": ""},
  234. "child_state2": {"value": ""},
  235. "complex": {
  236. 1: {"prop1": 42, "prop2": "hello"},
  237. 2: {"prop1": 42, "prop2": "hello"},
  238. },
  239. "dt": "1989-11-09 18:53:00+01:00",
  240. "fig": [],
  241. "is_hydrated": False,
  242. "key": "",
  243. "map_key": "a",
  244. "mapping": {"a": [1, 2, 3], "b": [4, 5, 6]},
  245. "num1": 0,
  246. "num2": 3.14,
  247. "obj": {"prop1": 42, "prop2": "hello"},
  248. "sum": 3.14,
  249. "upper": "",
  250. }
  251. assert formatted_state == exp_formatted_state
  252. def test_format_state_datetime():
  253. """Test that the format state is correct for datetime classes."""
  254. class DateTimeState(State):
  255. d: datetime.date = datetime.date.fromisoformat("1989-11-09")
  256. dt: datetime.datetime = datetime.datetime.fromisoformat(
  257. "1989-11-09T18:53:00+01:00"
  258. )
  259. t: datetime.time = datetime.time.fromisoformat("18:53:00+01:00")
  260. td: datetime.timedelta = datetime.timedelta(days=11, minutes=11)
  261. formatted_state = format.format_state(DateTimeState().dict())
  262. exp_formatted_state = {
  263. "d": "1989-11-09",
  264. "dt": "1989-11-09 18:53:00+01:00",
  265. "is_hydrated": False,
  266. "t": "18:53:00+01:00",
  267. "td": "11 days, 0:11:00",
  268. }
  269. assert formatted_state == exp_formatted_state
  270. def test_default_setters(test_state):
  271. """Test that we can set default values.
  272. Args:
  273. test_state: A state.
  274. """
  275. for prop_name in test_state.base_vars:
  276. # Each base var should have a default setter.
  277. assert hasattr(test_state, f"set_{prop_name}")
  278. def test_class_indexing_with_vars():
  279. """Test that we can index into a state var with another var."""
  280. prop = TestState.array[TestState.num1]
  281. assert str(prop) == "{test_state.array.at(test_state.num1)}"
  282. prop = TestState.mapping["a"][TestState.num1]
  283. assert str(prop) == '{test_state.mapping["a"].at(test_state.num1)}'
  284. prop = TestState.mapping[TestState.map_key]
  285. assert str(prop) == "{test_state.mapping[test_state.map_key]}"
  286. def test_class_attributes():
  287. """Test that we can get class attributes."""
  288. prop = TestState.obj.prop1
  289. assert str(prop) == "{test_state.obj.prop1}"
  290. prop = TestState.complex[1].prop1
  291. assert str(prop) == "{test_state.complex[1].prop1}"
  292. def test_get_parent_state():
  293. """Test getting the parent state."""
  294. assert TestState.get_parent_state() is None
  295. assert ChildState.get_parent_state() == TestState
  296. assert ChildState2.get_parent_state() == TestState
  297. assert GrandchildState.get_parent_state() == ChildState
  298. def test_get_substates():
  299. """Test getting the substates."""
  300. assert TestState.get_substates() == {ChildState, ChildState2}
  301. assert ChildState.get_substates() == {GrandchildState}
  302. assert ChildState2.get_substates() == set()
  303. assert GrandchildState.get_substates() == set()
  304. def test_get_name():
  305. """Test getting the name of a state."""
  306. assert TestState.get_name() == "test_state"
  307. assert ChildState.get_name() == "child_state"
  308. assert ChildState2.get_name() == "child_state2"
  309. assert GrandchildState.get_name() == "grandchild_state"
  310. def test_get_full_name():
  311. """Test getting the full name."""
  312. assert TestState.get_full_name() == "test_state"
  313. assert ChildState.get_full_name() == "test_state.child_state"
  314. assert ChildState2.get_full_name() == "test_state.child_state2"
  315. assert GrandchildState.get_full_name() == "test_state.child_state.grandchild_state"
  316. def test_get_class_substate():
  317. """Test getting the substate of a class."""
  318. assert TestState.get_class_substate(("child_state",)) == ChildState
  319. assert TestState.get_class_substate(("child_state2",)) == ChildState2
  320. assert ChildState.get_class_substate(("grandchild_state",)) == GrandchildState
  321. assert (
  322. TestState.get_class_substate(("child_state", "grandchild_state"))
  323. == GrandchildState
  324. )
  325. with pytest.raises(ValueError):
  326. TestState.get_class_substate(("invalid_child",))
  327. with pytest.raises(ValueError):
  328. TestState.get_class_substate(
  329. (
  330. "child_state",
  331. "invalid_child",
  332. )
  333. )
  334. def test_get_class_var():
  335. """Test getting the var of a class."""
  336. assert TestState.get_class_var(("num1",)).equals(TestState.num1)
  337. assert TestState.get_class_var(("num2",)).equals(TestState.num2)
  338. assert ChildState.get_class_var(("value",)).equals(ChildState.value)
  339. assert GrandchildState.get_class_var(("value2",)).equals(GrandchildState.value2)
  340. assert TestState.get_class_var(("child_state", "value")).equals(ChildState.value)
  341. assert TestState.get_class_var(
  342. ("child_state", "grandchild_state", "value2")
  343. ).equals(
  344. GrandchildState.value2,
  345. )
  346. assert ChildState.get_class_var(("grandchild_state", "value2")).equals(
  347. GrandchildState.value2,
  348. )
  349. with pytest.raises(ValueError):
  350. TestState.get_class_var(("invalid_var",))
  351. with pytest.raises(ValueError):
  352. TestState.get_class_var(
  353. (
  354. "child_state",
  355. "invalid_var",
  356. )
  357. )
  358. def test_set_class_var():
  359. """Test setting the var of a class."""
  360. with pytest.raises(AttributeError):
  361. TestState.num3 # type: ignore
  362. TestState._set_var(BaseVar(name="num3", type_=int).set_state(TestState))
  363. var = TestState.num3 # type: ignore
  364. assert var.name == "num3"
  365. assert var.type_ == int
  366. assert var.state == TestState.get_full_name()
  367. def test_set_parent_and_substates(test_state, child_state, grandchild_state):
  368. """Test setting the parent and substates.
  369. Args:
  370. test_state: A state.
  371. child_state: A child state.
  372. grandchild_state: A grandchild state.
  373. """
  374. assert len(test_state.substates) == 2
  375. assert set(test_state.substates) == {"child_state", "child_state2"}
  376. assert child_state.parent_state == test_state
  377. assert len(child_state.substates) == 1
  378. assert set(child_state.substates) == {"grandchild_state"}
  379. assert grandchild_state.parent_state == child_state
  380. assert len(grandchild_state.substates) == 0
  381. def test_get_child_attribute(test_state, child_state, child_state2, grandchild_state):
  382. """Test getting the attribute of a state.
  383. Args:
  384. test_state: A state.
  385. child_state: A child state.
  386. child_state2: A child state.
  387. grandchild_state: A grandchild state.
  388. """
  389. assert test_state.num1 == 0
  390. assert child_state.value == ""
  391. assert child_state2.value == ""
  392. assert child_state.count == 23
  393. assert grandchild_state.value2 == ""
  394. with pytest.raises(AttributeError):
  395. test_state.invalid
  396. with pytest.raises(AttributeError):
  397. test_state.child_state.invalid
  398. with pytest.raises(AttributeError):
  399. test_state.child_state.grandchild_state.invalid
  400. def test_set_child_attribute(test_state, child_state, grandchild_state):
  401. """Test setting the attribute of a state.
  402. Args:
  403. test_state: A state.
  404. child_state: A child state.
  405. grandchild_state: A grandchild state.
  406. """
  407. test_state.num1 = 10
  408. assert test_state.num1 == 10
  409. assert child_state.num1 == 10
  410. assert grandchild_state.num1 == 10
  411. grandchild_state.num1 = 5
  412. assert test_state.num1 == 5
  413. assert child_state.num1 == 5
  414. assert grandchild_state.num1 == 5
  415. child_state.value = "test"
  416. assert child_state.value == "test"
  417. assert grandchild_state.value == "test"
  418. grandchild_state.value = "test2"
  419. assert child_state.value == "test2"
  420. assert grandchild_state.value == "test2"
  421. grandchild_state.value2 = "test3"
  422. assert grandchild_state.value2 == "test3"
  423. def test_get_substate(test_state, child_state, child_state2, grandchild_state):
  424. """Test getting the substate of a state.
  425. Args:
  426. test_state: A state.
  427. child_state: A child state.
  428. child_state2: A child state.
  429. grandchild_state: A grandchild state.
  430. """
  431. assert test_state.get_substate(("child_state",)) == child_state
  432. assert test_state.get_substate(("child_state2",)) == child_state2
  433. assert (
  434. test_state.get_substate(("child_state", "grandchild_state")) == grandchild_state
  435. )
  436. assert child_state.get_substate(("grandchild_state",)) == grandchild_state
  437. with pytest.raises(ValueError):
  438. test_state.get_substate(("invalid",))
  439. with pytest.raises(ValueError):
  440. test_state.get_substate(("child_state", "invalid"))
  441. with pytest.raises(ValueError):
  442. test_state.get_substate(("child_state", "grandchild_state", "invalid"))
  443. def test_set_dirty_var(test_state):
  444. """Test changing state vars marks the value as dirty.
  445. Args:
  446. test_state: A state.
  447. """
  448. # Initially there should be no dirty vars.
  449. assert test_state.dirty_vars == set()
  450. # Setting a var should mark it as dirty.
  451. test_state.num1 = 1
  452. assert test_state.dirty_vars == {"num1", "sum"}
  453. # Setting another var should mark it as dirty.
  454. test_state.num2 = 2
  455. assert test_state.dirty_vars == {"num1", "num2", "sum"}
  456. # Cleaning the state should remove all dirty vars.
  457. test_state._clean()
  458. assert test_state.dirty_vars == set()
  459. def test_set_dirty_substate(test_state, child_state, child_state2, grandchild_state):
  460. """Test changing substate vars marks the value as dirty.
  461. Args:
  462. test_state: A state.
  463. child_state: A child state.
  464. child_state2: A child state.
  465. grandchild_state: A grandchild state.
  466. """
  467. # Initially there should be no dirty vars.
  468. assert test_state.dirty_vars == set()
  469. assert child_state.dirty_vars == set()
  470. assert child_state2.dirty_vars == set()
  471. assert grandchild_state.dirty_vars == set()
  472. # Setting a var should mark it as dirty.
  473. child_state.value = "test"
  474. assert child_state.dirty_vars == {"value"}
  475. assert test_state.dirty_substates == {"child_state"}
  476. assert child_state.dirty_substates == set()
  477. # Cleaning the parent state should remove the dirty substate.
  478. test_state._clean()
  479. assert test_state.dirty_substates == set()
  480. assert child_state.dirty_vars == set()
  481. # Setting a var on the grandchild should bubble up.
  482. grandchild_state.value2 = "test2"
  483. assert child_state.dirty_substates == {"grandchild_state"}
  484. assert test_state.dirty_substates == {"child_state"}
  485. # Cleaning the middle state should keep the parent state dirty.
  486. child_state._clean()
  487. assert test_state.dirty_substates == {"child_state"}
  488. assert child_state.dirty_substates == set()
  489. assert grandchild_state.dirty_vars == set()
  490. def test_reset(test_state, child_state):
  491. """Test resetting the state.
  492. Args:
  493. test_state: A state.
  494. child_state: A child state.
  495. """
  496. # Set some values.
  497. test_state.num1 = 1
  498. test_state.num2 = 2
  499. child_state.value = "test"
  500. # Reset the state.
  501. test_state.reset()
  502. # The values should be reset.
  503. assert test_state.num1 == 0
  504. assert test_state.num2 == 3.14
  505. assert child_state.value == ""
  506. expected_dirty_vars = {
  507. "num1",
  508. "num2",
  509. "obj",
  510. "upper",
  511. "complex",
  512. "is_hydrated",
  513. "fig",
  514. "key",
  515. "sum",
  516. "array",
  517. "map_key",
  518. "mapping",
  519. "dt",
  520. }
  521. # The dirty vars should be reset.
  522. assert test_state.dirty_vars == expected_dirty_vars
  523. assert child_state.dirty_vars == {"count", "value"}
  524. # The dirty substates should be reset.
  525. assert test_state.dirty_substates == {"child_state", "child_state2"}
  526. @pytest.mark.asyncio
  527. async def test_process_event_simple(test_state):
  528. """Test processing an event.
  529. Args:
  530. test_state: A state.
  531. """
  532. assert test_state.num1 == 0
  533. event = Event(token="t", name="set_num1", payload={"value": 69})
  534. update = await test_state._process(event).__anext__()
  535. # The event should update the value.
  536. assert test_state.num1 == 69
  537. # The delta should contain the changes, including computed vars.
  538. # assert update.delta == {"test_state": {"num1": 69, "sum": 72.14}}
  539. assert update.delta == {"test_state": {"num1": 69, "sum": 72.14, "upper": ""}}
  540. assert update.events == []
  541. @pytest.mark.asyncio
  542. async def test_process_event_substate(test_state, child_state, grandchild_state):
  543. """Test processing an event on a substate.
  544. Args:
  545. test_state: A state.
  546. child_state: A child state.
  547. grandchild_state: A grandchild state.
  548. """
  549. # Events should bubble down to the substate.
  550. assert child_state.value == ""
  551. assert child_state.count == 23
  552. event = Event(
  553. token="t", name="child_state.change_both", payload={"value": "hi", "count": 12}
  554. )
  555. update = await test_state._process(event).__anext__()
  556. assert child_state.value == "HI"
  557. assert child_state.count == 24
  558. assert update.delta == {
  559. "test_state": {"sum": 3.14, "upper": ""},
  560. "test_state.child_state": {"value": "HI", "count": 24},
  561. }
  562. test_state._clean()
  563. # Test with the granchild state.
  564. assert grandchild_state.value2 == ""
  565. event = Event(
  566. token="t",
  567. name="child_state.grandchild_state.set_value2",
  568. payload={"value": "new"},
  569. )
  570. update = await test_state._process(event).__anext__()
  571. assert grandchild_state.value2 == "new"
  572. assert update.delta == {
  573. "test_state": {"sum": 3.14, "upper": ""},
  574. "test_state.child_state.grandchild_state": {"value2": "new"},
  575. }
  576. @pytest.mark.asyncio
  577. async def test_process_event_generator():
  578. """Test event handlers that generate multiple updates."""
  579. gen_state = GenState() # type: ignore
  580. event = Event(
  581. token="t",
  582. name="go",
  583. payload={"c": 5},
  584. )
  585. gen = gen_state._process(event)
  586. count = 0
  587. async for update in gen:
  588. count += 1
  589. if count == 6:
  590. assert update.delta == {}
  591. assert update.final
  592. else:
  593. assert gen_state.value == count
  594. assert update.delta == {
  595. "gen_state": {"value": count},
  596. }
  597. assert not update.final
  598. assert count == 6
  599. def test_format_event_handler():
  600. """Test formatting an event handler."""
  601. assert (
  602. format.format_event_handler(TestState.do_something) == "test_state.do_something" # type: ignore
  603. )
  604. assert (
  605. format.format_event_handler(ChildState.change_both) # type: ignore
  606. == "test_state.child_state.change_both"
  607. )
  608. assert (
  609. format.format_event_handler(GrandchildState.do_nothing) # type: ignore
  610. == "test_state.child_state.grandchild_state.do_nothing"
  611. )
  612. def test_get_token(test_state, mocker, router_data):
  613. """Test that the token obtained from the router_data is correct.
  614. Args:
  615. test_state: The test state.
  616. mocker: Pytest Mocker object.
  617. router_data: The router data fixture.
  618. """
  619. mocker.patch.object(test_state, "router_data", router_data)
  620. assert test_state.get_token() == "b181904c-3953-4a79-dc18-ae9518c22f05"
  621. def test_get_sid(test_state, mocker, router_data):
  622. """Test getting session id.
  623. Args:
  624. test_state: A state.
  625. mocker: Pytest Mocker object.
  626. router_data: The router data fixture.
  627. """
  628. mocker.patch.object(test_state, "router_data", router_data)
  629. assert test_state.get_sid() == "9fpxSzPb9aFMb4wFAAAH"
  630. def test_get_headers(test_state, mocker, router_data, router_data_headers):
  631. """Test getting client headers.
  632. Args:
  633. test_state: A state.
  634. mocker: Pytest Mocker object.
  635. router_data: The router data fixture.
  636. router_data_headers: The expected headers.
  637. """
  638. mocker.patch.object(test_state, "router_data", router_data)
  639. assert test_state.get_headers() == router_data_headers
  640. def test_get_client_ip(test_state, mocker, router_data):
  641. """Test getting client IP.
  642. Args:
  643. test_state: A state.
  644. mocker: Pytest Mocker object.
  645. router_data: The router data fixture.
  646. """
  647. mocker.patch.object(test_state, "router_data", router_data)
  648. assert test_state.get_client_ip() == "127.0.0.1"
  649. def test_get_cookies(test_state, mocker, router_data):
  650. """Test getting client cookies.
  651. Args:
  652. test_state: A state.
  653. mocker: Pytest Mocker object.
  654. router_data: The router data fixture.
  655. """
  656. mocker.patch.object(test_state, "router_data", router_data)
  657. assert test_state.get_cookies() == {
  658. "csrftoken": "mocktoken",
  659. "name": "reflex",
  660. "list_cookies": ["some", "random", "cookies"],
  661. "dict_cookies": {"name": "reflex"},
  662. "val": True,
  663. }
  664. def test_get_current_page(test_state):
  665. assert test_state.get_current_page() == ""
  666. route = "mypage/subpage"
  667. test_state.router_data = {RouteVar.PATH: route}
  668. assert test_state.get_current_page() == route
  669. def test_get_query_params(test_state):
  670. assert test_state.get_query_params() == {}
  671. params = {"p1": "a", "p2": "b"}
  672. test_state.router_data = {RouteVar.QUERY: params}
  673. assert test_state.get_query_params() == params
  674. def test_add_var():
  675. class DynamicState(State):
  676. pass
  677. ds1 = DynamicState()
  678. assert "dynamic_int" not in ds1.__dict__
  679. assert not hasattr(ds1, "dynamic_int")
  680. ds1.add_var("dynamic_int", int, 42)
  681. # Existing instances get the BaseVar
  682. assert ds1.dynamic_int.equals(DynamicState.dynamic_int) # type: ignore
  683. # New instances get an actual value with the default
  684. assert DynamicState().dynamic_int == 42
  685. ds1.add_var("dynamic_list", List[int], [5, 10])
  686. assert ds1.dynamic_list.equals(DynamicState.dynamic_list) # type: ignore
  687. ds2 = DynamicState()
  688. assert ds2.dynamic_list == [5, 10]
  689. ds2.dynamic_list.append(15)
  690. assert ds2.dynamic_list == [5, 10, 15]
  691. assert DynamicState().dynamic_list == [5, 10]
  692. ds1.add_var("dynamic_dict", Dict[str, int], {"k1": 5, "k2": 10})
  693. assert ds1.dynamic_dict.equals(DynamicState.dynamic_dict) # type: ignore
  694. assert ds2.dynamic_dict.equals(DynamicState.dynamic_dict) # type: ignore
  695. assert DynamicState().dynamic_dict == {"k1": 5, "k2": 10}
  696. assert DynamicState().dynamic_dict == {"k1": 5, "k2": 10}
  697. def test_add_var_default_handlers(test_state):
  698. test_state.add_var("rand_int", int, 10)
  699. assert "set_rand_int" in test_state.event_handlers
  700. assert isinstance(test_state.event_handlers["set_rand_int"], EventHandler)
  701. class InterdependentState(State):
  702. """A state with 3 vars and 3 computed vars.
  703. x: a variable that no computed var depends on
  704. v1: a varable that one computed var directly depeneds on
  705. _v2: a backend variable that one computed var directly depends on
  706. v1x2: a computed var that depends on v1
  707. v2x2: a computed var that depends on backend var _v2
  708. v1x2x2: a computed var that depends on computed var v1x2
  709. """
  710. x: int = 0
  711. v1: int = 0
  712. _v2: int = 1
  713. @rx.cached_var
  714. def v1x2(self) -> int:
  715. """Depends on var v1.
  716. Returns:
  717. Var v1 multiplied by 2
  718. """
  719. return self.v1 * 2
  720. @rx.cached_var
  721. def v2x2(self) -> int:
  722. """Depends on backend var _v2.
  723. Returns:
  724. backend var _v2 multiplied by 2
  725. """
  726. return self._v2 * 2
  727. @rx.cached_var
  728. def v1x2x2(self) -> int:
  729. """Depends on ComputedVar v1x2.
  730. Returns:
  731. ComputedVar v1x2 multiplied by 2
  732. """
  733. return self.v1x2 * 2
  734. @pytest.fixture
  735. def interdependent_state() -> State:
  736. """A state with varying dependency between vars.
  737. Returns:
  738. instance of InterdependentState
  739. """
  740. s = InterdependentState()
  741. s.dict() # prime initial relationships by accessing all ComputedVars
  742. return s
  743. def test_not_dirty_computed_var_from_var(interdependent_state):
  744. """Set Var that no ComputedVar depends on, expect no recalculation.
  745. Args:
  746. interdependent_state: A state with varying Var dependencies.
  747. """
  748. interdependent_state.x = 5
  749. assert interdependent_state.get_delta() == {
  750. interdependent_state.get_full_name(): {"x": 5},
  751. }
  752. def test_dirty_computed_var_from_var(interdependent_state):
  753. """Set Var that ComputedVar depends on, expect recalculation.
  754. The other ComputedVar depends on the changed ComputedVar and should also be
  755. recalculated. No other ComputedVars should be recalculated.
  756. Args:
  757. interdependent_state: A state with varying Var dependencies.
  758. """
  759. interdependent_state.v1 = 1
  760. assert interdependent_state.get_delta() == {
  761. interdependent_state.get_full_name(): {"v1": 1, "v1x2": 2, "v1x2x2": 4},
  762. }
  763. def test_dirty_computed_var_from_backend_var(interdependent_state):
  764. """Set backend var that ComputedVar depends on, expect recalculation.
  765. Args:
  766. interdependent_state: A state with varying Var dependencies.
  767. """
  768. interdependent_state._v2 = 2
  769. assert interdependent_state.get_delta() == {
  770. interdependent_state.get_full_name(): {"v2x2": 4},
  771. }
  772. def test_per_state_backend_var(interdependent_state):
  773. """Set backend var on one instance, expect no affect in other instances.
  774. Args:
  775. interdependent_state: A state with varying Var dependencies.
  776. """
  777. s2 = InterdependentState()
  778. assert s2._v2 == interdependent_state._v2
  779. interdependent_state._v2 = 2
  780. assert s2._v2 != interdependent_state._v2
  781. s3 = InterdependentState()
  782. assert s3._v2 != interdependent_state._v2
  783. # both s2 and s3 should still have the default value
  784. assert s2._v2 == s3._v2
  785. # changing s2._v2 should not affect others
  786. s2._v2 = 4
  787. assert s2._v2 != interdependent_state._v2
  788. assert s2._v2 != s3._v2
  789. def test_child_state():
  790. """Test that the child state computed vars can reference parent state vars."""
  791. class MainState(State):
  792. v: int = 2
  793. class ChildState(MainState):
  794. @ComputedVar
  795. def rendered_var(self):
  796. return self.v
  797. ms = MainState()
  798. cs = ms.substates[ChildState.get_name()]
  799. assert ms.v == 2
  800. assert cs.v == 2
  801. assert cs.rendered_var == 2
  802. def test_conditional_computed_vars():
  803. """Test that computed vars can have conditionals."""
  804. class MainState(State):
  805. flag: bool = False
  806. t1: str = "a"
  807. t2: str = "b"
  808. @ComputedVar
  809. def rendered_var(self) -> str:
  810. if self.flag:
  811. return self.t1
  812. return self.t2
  813. ms = MainState()
  814. # Initially there are no dirty computed vars.
  815. assert ms._dirty_computed_vars(from_vars={"flag"}) == {"rendered_var"}
  816. assert ms._dirty_computed_vars(from_vars={"t2"}) == {"rendered_var"}
  817. assert ms._dirty_computed_vars(from_vars={"t1"}) == {"rendered_var"}
  818. assert ms.computed_vars["rendered_var"].deps(objclass=MainState) == {
  819. "flag",
  820. "t1",
  821. "t2",
  822. }
  823. def test_event_handlers_convert_to_fns(test_state, child_state):
  824. """Test that when the state is initialized, event handlers are converted to fns.
  825. Args:
  826. test_state: A state with event handlers.
  827. child_state: A child state with event handlers.
  828. """
  829. # The class instances should be event handlers.
  830. assert isinstance(TestState.do_something, EventHandler)
  831. assert isinstance(ChildState.change_both, EventHandler)
  832. # The object instances should be fns.
  833. test_state.do_something()
  834. child_state.change_both(value="goose", count=9)
  835. assert child_state.value == "GOOSE"
  836. assert child_state.count == 18
  837. def test_event_handlers_call_other_handlers():
  838. """Test that event handlers can call other event handlers."""
  839. class MainState(State):
  840. v: int = 0
  841. def set_v(self, v: int):
  842. self.v = v
  843. def set_v2(self, v: int):
  844. self.set_v(v)
  845. class SubState(MainState):
  846. def set_v3(self, v: int):
  847. self.set_v2(v)
  848. ms = MainState()
  849. ms.set_v2(1)
  850. assert ms.v == 1
  851. # ensure handler can be called from substate
  852. ms.substates[SubState.get_name()].set_v3(2)
  853. assert ms.v == 2
  854. def test_computed_var_cached():
  855. """Test that a ComputedVar doesn't recalculate when accessed."""
  856. comp_v_calls = 0
  857. class ComputedState(State):
  858. v: int = 0
  859. @rx.cached_var
  860. def comp_v(self) -> int:
  861. nonlocal comp_v_calls
  862. comp_v_calls += 1
  863. return self.v
  864. cs = ComputedState()
  865. assert cs.dict()["v"] == 0
  866. assert comp_v_calls == 1
  867. assert cs.dict()["comp_v"] == 0
  868. assert comp_v_calls == 1
  869. assert cs.comp_v == 0
  870. assert comp_v_calls == 1
  871. cs.v = 1
  872. assert comp_v_calls == 1
  873. assert cs.comp_v == 1
  874. assert comp_v_calls == 2
  875. def test_computed_var_cached_depends_on_non_cached():
  876. """Test that a cached_var is recalculated if it depends on non-cached ComputedVar."""
  877. class ComputedState(State):
  878. v: int = 0
  879. @rx.var
  880. def no_cache_v(self) -> int:
  881. return self.v
  882. @rx.cached_var
  883. def dep_v(self) -> int:
  884. return self.no_cache_v
  885. @rx.cached_var
  886. def comp_v(self) -> int:
  887. return self.v
  888. cs = ComputedState()
  889. assert cs.dirty_vars == set()
  890. assert cs.get_delta() == {cs.get_name(): {"no_cache_v": 0, "dep_v": 0}}
  891. cs._clean()
  892. assert cs.dirty_vars == set()
  893. assert cs.get_delta() == {cs.get_name(): {"no_cache_v": 0, "dep_v": 0}}
  894. cs._clean()
  895. assert cs.dirty_vars == set()
  896. cs.v = 1
  897. assert cs.dirty_vars == {"v", "comp_v", "dep_v", "no_cache_v"}
  898. assert cs.get_delta() == {
  899. cs.get_name(): {"v": 1, "no_cache_v": 1, "dep_v": 1, "comp_v": 1}
  900. }
  901. cs._clean()
  902. assert cs.dirty_vars == set()
  903. assert cs.get_delta() == {cs.get_name(): {"no_cache_v": 1, "dep_v": 1}}
  904. cs._clean()
  905. assert cs.dirty_vars == set()
  906. assert cs.get_delta() == {cs.get_name(): {"no_cache_v": 1, "dep_v": 1}}
  907. cs._clean()
  908. assert cs.dirty_vars == set()
  909. def test_computed_var_depends_on_parent_non_cached():
  910. """Child state cached_var that depends on parent state un cached var is always recalculated."""
  911. counter = 0
  912. class ParentState(State):
  913. @rx.var
  914. def no_cache_v(self) -> int:
  915. nonlocal counter
  916. counter += 1
  917. return counter
  918. class ChildState(ParentState):
  919. @rx.cached_var
  920. def dep_v(self) -> int:
  921. return self.no_cache_v
  922. ps = ParentState()
  923. cs = ps.substates[ChildState.get_name()]
  924. assert ps.dirty_vars == set()
  925. assert cs.dirty_vars == set()
  926. assert ps.dict() == {
  927. cs.get_name(): {"dep_v": 2},
  928. "no_cache_v": 1,
  929. IS_HYDRATED: False,
  930. }
  931. assert ps.dict() == {
  932. cs.get_name(): {"dep_v": 4},
  933. "no_cache_v": 3,
  934. IS_HYDRATED: False,
  935. }
  936. assert ps.dict() == {
  937. cs.get_name(): {"dep_v": 6},
  938. "no_cache_v": 5,
  939. IS_HYDRATED: False,
  940. }
  941. assert counter == 6
  942. @pytest.mark.parametrize("use_partial", [True, False])
  943. def test_cached_var_depends_on_event_handler(use_partial: bool):
  944. """A cached_var that calls an event handler calculates deps correctly.
  945. Args:
  946. use_partial: if true, replace the EventHandler with functools.partial
  947. """
  948. counter = 0
  949. class HandlerState(State):
  950. x: int = 42
  951. def handler(self):
  952. self.x = self.x + 1
  953. @rx.cached_var
  954. def cached_x_side_effect(self) -> int:
  955. self.handler()
  956. nonlocal counter
  957. counter += 1
  958. return counter
  959. if use_partial:
  960. HandlerState.handler = functools.partial(HandlerState.handler.fn)
  961. assert isinstance(HandlerState.handler, functools.partial)
  962. else:
  963. assert isinstance(HandlerState.handler, EventHandler)
  964. s = HandlerState()
  965. assert "cached_x_side_effect" in s.computed_var_dependencies["x"]
  966. assert s.cached_x_side_effect == 1
  967. assert s.x == 43
  968. s.handler()
  969. assert s.cached_x_side_effect == 2
  970. assert s.x == 45
  971. def test_computed_var_dependencies():
  972. """Test that a ComputedVar correctly tracks its dependencies."""
  973. class ComputedState(State):
  974. v: int = 0
  975. w: int = 0
  976. x: int = 0
  977. y: List[int] = [1, 2, 3]
  978. _z: List[int] = [1, 2, 3]
  979. @rx.cached_var
  980. def comp_v(self) -> int:
  981. """Direct access.
  982. Returns:
  983. The value of self.v.
  984. """
  985. return self.v
  986. @rx.cached_var
  987. def comp_w(self):
  988. """Nested lambda.
  989. Returns:
  990. A lambda that returns the value of self.w.
  991. """
  992. return lambda: self.w
  993. @rx.cached_var
  994. def comp_x(self):
  995. """Nested function.
  996. Returns:
  997. A function that returns the value of self.x.
  998. """
  999. def _():
  1000. return self.x
  1001. return _
  1002. @rx.cached_var
  1003. def comp_y(self) -> List[int]:
  1004. """Comprehension iterating over attribute.
  1005. Returns:
  1006. A list of the values of self.y.
  1007. """
  1008. return [round(y) for y in self.y]
  1009. @rx.cached_var
  1010. def comp_z(self) -> List[bool]:
  1011. """Comprehension accesses attribute.
  1012. Returns:
  1013. A list of whether the values 0-4 are in self._z.
  1014. """
  1015. return [z in self._z for z in range(5)]
  1016. cs = ComputedState()
  1017. assert cs.computed_var_dependencies["v"] == {"comp_v"}
  1018. assert cs.computed_var_dependencies["w"] == {"comp_w"}
  1019. assert cs.computed_var_dependencies["x"] == {"comp_x"}
  1020. assert cs.computed_var_dependencies["y"] == {"comp_y"}
  1021. assert cs.computed_var_dependencies["_z"] == {"comp_z"}
  1022. def test_backend_method():
  1023. """A method with leading underscore should be callable from event handler."""
  1024. class BackendMethodState(State):
  1025. def _be_method(self):
  1026. return True
  1027. def handler(self):
  1028. assert self._be_method()
  1029. bms = BackendMethodState()
  1030. bms.handler()
  1031. assert bms._be_method()
  1032. def test_setattr_of_mutable_types(mutable_state):
  1033. """Test that mutable types are converted to corresponding Reflex wrappers.
  1034. Args:
  1035. mutable_state: A test state.
  1036. """
  1037. array = mutable_state.array
  1038. hashmap = mutable_state.hashmap
  1039. test_set = mutable_state.test_set
  1040. assert isinstance(array, MutableProxy)
  1041. assert isinstance(array, list)
  1042. assert isinstance(array[1], MutableProxy)
  1043. assert isinstance(array[1], list)
  1044. assert isinstance(array[2], MutableProxy)
  1045. assert isinstance(array[2], dict)
  1046. assert isinstance(hashmap, MutableProxy)
  1047. assert isinstance(hashmap, dict)
  1048. assert isinstance(hashmap["key"], MutableProxy)
  1049. assert isinstance(hashmap["key"], list)
  1050. assert isinstance(hashmap["third_key"], MutableProxy)
  1051. assert isinstance(hashmap["third_key"], dict)
  1052. assert isinstance(test_set, MutableProxy)
  1053. assert isinstance(test_set, set)
  1054. assert isinstance(mutable_state.custom, MutableProxy)
  1055. assert isinstance(mutable_state.custom.array, MutableProxy)
  1056. assert isinstance(mutable_state.custom.array, list)
  1057. assert isinstance(mutable_state.custom.hashmap, MutableProxy)
  1058. assert isinstance(mutable_state.custom.hashmap, dict)
  1059. assert isinstance(mutable_state.custom.test_set, MutableProxy)
  1060. assert isinstance(mutable_state.custom.test_set, set)
  1061. assert isinstance(mutable_state.custom.custom, MutableProxy)
  1062. mutable_state.reassign_mutables()
  1063. array = mutable_state.array
  1064. hashmap = mutable_state.hashmap
  1065. test_set = mutable_state.test_set
  1066. assert isinstance(array, MutableProxy)
  1067. assert isinstance(array, list)
  1068. assert isinstance(array[1], MutableProxy)
  1069. assert isinstance(array[1], list)
  1070. assert isinstance(array[2], MutableProxy)
  1071. assert isinstance(array[2], dict)
  1072. assert isinstance(hashmap, MutableProxy)
  1073. assert isinstance(hashmap, dict)
  1074. assert isinstance(hashmap["mod_key"], MutableProxy)
  1075. assert isinstance(hashmap["mod_key"], list)
  1076. assert isinstance(hashmap["mod_third_key"], MutableProxy)
  1077. assert isinstance(hashmap["mod_third_key"], dict)
  1078. assert isinstance(test_set, MutableProxy)
  1079. assert isinstance(test_set, set)
  1080. def test_error_on_state_method_shadow():
  1081. """Test that an error is thrown when an event handler shadows a state method."""
  1082. with pytest.raises(NameError) as err:
  1083. class InvalidTest(rx.State):
  1084. def reset(self):
  1085. pass
  1086. assert (
  1087. err.value.args[0]
  1088. == f"The event handler name `reset` shadows a builtin State method; use a different name instead"
  1089. )
  1090. def test_state_with_invalid_yield():
  1091. """Test that an error is thrown when a state yields an invalid value."""
  1092. class StateWithInvalidYield(rx.State):
  1093. """A state that yields an invalid value."""
  1094. def invalid_handler(self):
  1095. """Invalid handler.
  1096. Yields:
  1097. an invalid value.
  1098. """
  1099. yield 1
  1100. invalid_state = StateWithInvalidYield()
  1101. with pytest.raises(TypeError) as err:
  1102. invalid_state._check_valid(
  1103. invalid_state.event_handlers["invalid_handler"],
  1104. rx.event.Event(token="fake_token", name="invalid_handler"),
  1105. )
  1106. assert (
  1107. "must only return/yield: None, Events or other EventHandlers"
  1108. in err.value.args[0]
  1109. )
  1110. @pytest.fixture(scope="function", params=["in_process", "redis"])
  1111. def state_manager(request) -> Generator[StateManager, None, None]:
  1112. """Instance of state manager parametrized for redis and in-process.
  1113. Args:
  1114. request: pytest request object.
  1115. Yields:
  1116. A state manager instance
  1117. """
  1118. state_manager = StateManager.create(state=TestState)
  1119. if request.param == "redis":
  1120. if not isinstance(state_manager, StateManagerRedis):
  1121. pytest.skip("Test requires redis")
  1122. else:
  1123. # explicitly NOT using redis
  1124. state_manager = StateManagerMemory(state=TestState)
  1125. assert not state_manager._states_locks
  1126. yield state_manager
  1127. if isinstance(state_manager, StateManagerRedis):
  1128. asyncio.get_event_loop().run_until_complete(state_manager.redis.close())
  1129. @pytest.mark.asyncio
  1130. async def test_state_manager_modify_state(state_manager: StateManager, token: str):
  1131. """Test that the state manager can modify a state exclusively.
  1132. Args:
  1133. state_manager: A state manager instance.
  1134. token: A token.
  1135. """
  1136. async with state_manager.modify_state(token):
  1137. if isinstance(state_manager, StateManagerRedis):
  1138. assert await state_manager.redis.get(f"{token}_lock")
  1139. elif isinstance(state_manager, StateManagerMemory):
  1140. assert token in state_manager._states_locks
  1141. assert state_manager._states_locks[token].locked()
  1142. # lock should be dropped after exiting the context
  1143. if isinstance(state_manager, StateManagerRedis):
  1144. assert (await state_manager.redis.get(f"{token}_lock")) is None
  1145. elif isinstance(state_manager, StateManagerMemory):
  1146. assert not state_manager._states_locks[token].locked()
  1147. # separate instances should NOT share locks
  1148. sm2 = StateManagerMemory(state=TestState)
  1149. assert sm2._state_manager_lock is state_manager._state_manager_lock
  1150. assert not sm2._states_locks
  1151. if state_manager._states_locks:
  1152. assert sm2._states_locks != state_manager._states_locks
  1153. @pytest.mark.asyncio
  1154. async def test_state_manager_contend(state_manager: StateManager, token: str):
  1155. """Multiple coroutines attempting to access the same state.
  1156. Args:
  1157. state_manager: A state manager instance.
  1158. token: A token.
  1159. """
  1160. n_coroutines = 10
  1161. exp_num1 = 10
  1162. async with state_manager.modify_state(token) as state:
  1163. state.num1 = 0
  1164. async def _coro():
  1165. async with state_manager.modify_state(token) as state:
  1166. await asyncio.sleep(0.01)
  1167. state.num1 += 1
  1168. tasks = [asyncio.create_task(_coro()) for _ in range(n_coroutines)]
  1169. for f in asyncio.as_completed(tasks):
  1170. await f
  1171. assert (await state_manager.get_state(token)).num1 == exp_num1
  1172. if isinstance(state_manager, StateManagerRedis):
  1173. assert (await state_manager.redis.get(f"{token}_lock")) is None
  1174. elif isinstance(state_manager, StateManagerMemory):
  1175. assert token in state_manager._states_locks
  1176. assert not state_manager._states_locks[token].locked()
  1177. @pytest.fixture(scope="function")
  1178. def state_manager_redis() -> Generator[StateManager, None, None]:
  1179. """Instance of state manager for redis only.
  1180. Yields:
  1181. A state manager instance
  1182. """
  1183. state_manager = StateManager.create(TestState)
  1184. if not isinstance(state_manager, StateManagerRedis):
  1185. pytest.skip("Test requires redis")
  1186. yield state_manager
  1187. asyncio.get_event_loop().run_until_complete(state_manager.redis.close())
  1188. @pytest.mark.asyncio
  1189. async def test_state_manager_lock_expire(state_manager_redis: StateManager, token: str):
  1190. """Test that the state manager lock expires and raises exception exiting context.
  1191. Args:
  1192. state_manager_redis: A state manager instance.
  1193. token: A token.
  1194. """
  1195. state_manager_redis.lock_expiration = LOCK_EXPIRATION
  1196. async with state_manager_redis.modify_state(token):
  1197. await asyncio.sleep(0.01)
  1198. with pytest.raises(LockExpiredError):
  1199. async with state_manager_redis.modify_state(token):
  1200. await asyncio.sleep(LOCK_EXPIRE_SLEEP)
  1201. @pytest.mark.asyncio
  1202. async def test_state_manager_lock_expire_contend(
  1203. state_manager_redis: StateManager, token: str
  1204. ):
  1205. """Test that the state manager lock expires and queued waiters proceed.
  1206. Args:
  1207. state_manager_redis: A state manager instance.
  1208. token: A token.
  1209. """
  1210. exp_num1 = 4252
  1211. unexp_num1 = 666
  1212. state_manager_redis.lock_expiration = LOCK_EXPIRATION
  1213. order = []
  1214. async def _coro_blocker():
  1215. async with state_manager_redis.modify_state(token) as state:
  1216. order.append("blocker")
  1217. await asyncio.sleep(LOCK_EXPIRE_SLEEP)
  1218. state.num1 = unexp_num1
  1219. async def _coro_waiter():
  1220. while "blocker" not in order:
  1221. await asyncio.sleep(0.005)
  1222. async with state_manager_redis.modify_state(token) as state:
  1223. order.append("waiter")
  1224. assert state.num1 != unexp_num1
  1225. state.num1 = exp_num1
  1226. tasks = [
  1227. asyncio.create_task(_coro_blocker()),
  1228. asyncio.create_task(_coro_waiter()),
  1229. ]
  1230. with pytest.raises(LockExpiredError):
  1231. await tasks[0]
  1232. await tasks[1]
  1233. assert order == ["blocker", "waiter"]
  1234. assert (await state_manager_redis.get_state(token)).num1 == exp_num1
  1235. @pytest.fixture(scope="function")
  1236. def mock_app(monkeypatch, app: rx.App, state_manager: StateManager) -> rx.App:
  1237. """Mock app fixture.
  1238. Args:
  1239. monkeypatch: Pytest monkeypatch object.
  1240. app: An app.
  1241. state_manager: A state manager.
  1242. Returns:
  1243. The app, after mocking out prerequisites.get_app()
  1244. """
  1245. app_module = Mock()
  1246. setattr(app_module, APP_VAR, app)
  1247. app.state = TestState
  1248. app.state_manager = state_manager
  1249. assert app.event_namespace is not None
  1250. app.event_namespace.emit = AsyncMock()
  1251. monkeypatch.setattr(prerequisites, "get_app", lambda: app_module)
  1252. return app
  1253. @pytest.mark.asyncio
  1254. async def test_state_proxy(grandchild_state: GrandchildState, mock_app: rx.App):
  1255. """Test that the state proxy works.
  1256. Args:
  1257. grandchild_state: A grandchild state.
  1258. mock_app: An app that will be returned by `get_app()`
  1259. """
  1260. child_state = grandchild_state.parent_state
  1261. assert child_state is not None
  1262. parent_state = child_state.parent_state
  1263. assert parent_state is not None
  1264. if isinstance(mock_app.state_manager, StateManagerMemory):
  1265. mock_app.state_manager.states[parent_state.get_token()] = parent_state
  1266. sp = StateProxy(grandchild_state)
  1267. assert sp.__wrapped__ == grandchild_state
  1268. assert sp._self_substate_path == grandchild_state.get_full_name().split(".")
  1269. assert sp._self_app is mock_app
  1270. assert not sp._self_mutable
  1271. assert sp._self_actx is None
  1272. # cannot use normal contextmanager protocol
  1273. with pytest.raises(TypeError), sp:
  1274. pass
  1275. with pytest.raises(ImmutableStateError):
  1276. # cannot directly modify state proxy outside of async context
  1277. sp.value2 = 16
  1278. async with sp:
  1279. assert sp._self_actx is not None
  1280. assert sp._self_mutable # proxy is mutable inside context
  1281. if isinstance(mock_app.state_manager, StateManagerMemory):
  1282. # For in-process store, only one instance of the state exists
  1283. assert sp.__wrapped__ is grandchild_state
  1284. else:
  1285. # When redis is used, a new+updated instance is assigned to the proxy
  1286. assert sp.__wrapped__ is not grandchild_state
  1287. sp.value2 = 42
  1288. assert not sp._self_mutable # proxy is not mutable after exiting context
  1289. assert sp._self_actx is None
  1290. assert sp.value2 == 42
  1291. # Get the state from the state manager directly and check that the value is updated
  1292. gotten_state = await mock_app.state_manager.get_state(grandchild_state.get_token())
  1293. if isinstance(mock_app.state_manager, StateManagerMemory):
  1294. # For in-process store, only one instance of the state exists
  1295. assert gotten_state is parent_state
  1296. else:
  1297. assert gotten_state is not parent_state
  1298. gotten_grandchild_state = gotten_state.get_substate(sp._self_substate_path)
  1299. assert gotten_grandchild_state is not None
  1300. assert gotten_grandchild_state.value2 == 42
  1301. # ensure state update was emitted
  1302. assert mock_app.event_namespace is not None
  1303. mock_app.event_namespace.emit.assert_called_once()
  1304. mcall = mock_app.event_namespace.emit.mock_calls[0]
  1305. assert mcall.args[0] == str(SocketEvent.EVENT)
  1306. assert json.loads(mcall.args[1]) == StateUpdate(
  1307. delta={
  1308. parent_state.get_full_name(): {
  1309. "upper": "",
  1310. "sum": 3.14,
  1311. },
  1312. grandchild_state.get_full_name(): {
  1313. "value2": 42,
  1314. },
  1315. }
  1316. )
  1317. assert mcall.kwargs["to"] == grandchild_state.get_sid()
  1318. class BackgroundTaskState(State):
  1319. """A state with a background task."""
  1320. order: List[str] = []
  1321. dict_list: Dict[str, List[int]] = {"foo": []}
  1322. @rx.background
  1323. async def background_task(self):
  1324. """A background task that updates the state."""
  1325. async with self:
  1326. assert not self.order
  1327. self.order.append("background_task:start")
  1328. assert isinstance(self, StateProxy)
  1329. with pytest.raises(ImmutableStateError):
  1330. self.order.append("bad idea")
  1331. with pytest.raises(ImmutableStateError):
  1332. # Even nested access to mutables raises an exception.
  1333. self.dict_list["foo"].append(42)
  1334. with pytest.raises(ImmutableStateError):
  1335. # Direct calling another handler that modifies state raises an exception.
  1336. self.other()
  1337. with pytest.raises(ImmutableStateError):
  1338. # Calling other methods that modify state raises an exception.
  1339. self._private_method()
  1340. # wait for some other event to happen
  1341. while len(self.order) == 1:
  1342. await asyncio.sleep(0.01)
  1343. async with self:
  1344. pass # update proxy instance
  1345. async with self:
  1346. self.order.append("background_task:stop")
  1347. self.other() # direct calling event handlers works in context
  1348. self._private_method()
  1349. @rx.background
  1350. async def background_task_reset(self):
  1351. """A background task that resets the state."""
  1352. with pytest.raises(ImmutableStateError):
  1353. # Resetting the state should be explicitly blocked.
  1354. self.reset()
  1355. async with self:
  1356. self.order.append("foo")
  1357. self.reset()
  1358. assert not self.order
  1359. async with self:
  1360. self.order.append("reset")
  1361. @rx.background
  1362. async def background_task_generator(self):
  1363. """A background task generator that does nothing.
  1364. Yields:
  1365. None
  1366. """
  1367. yield
  1368. def other(self):
  1369. """Some other event that updates the state."""
  1370. self.order.append("other")
  1371. def _private_method(self):
  1372. """Some private method that updates the state."""
  1373. self.order.append("private")
  1374. async def bad_chain1(self):
  1375. """Test that a background task cannot be chained."""
  1376. await self.background_task()
  1377. async def bad_chain2(self):
  1378. """Test that a background task generator cannot be chained."""
  1379. async for _foo in self.background_task_generator():
  1380. pass
  1381. @pytest.mark.asyncio
  1382. async def test_background_task_no_block(mock_app: rx.App, token: str):
  1383. """Test that a background task does not block other events.
  1384. Args:
  1385. mock_app: An app that will be returned by `get_app()`
  1386. token: A token.
  1387. """
  1388. router_data = {"query": {}}
  1389. mock_app.state_manager.state = mock_app.state = BackgroundTaskState
  1390. async for update in rx.app.process( # type: ignore
  1391. mock_app,
  1392. Event(
  1393. token=token,
  1394. name=f"{BackgroundTaskState.get_name()}.background_task",
  1395. router_data=router_data,
  1396. payload={},
  1397. ),
  1398. sid="",
  1399. headers={},
  1400. client_ip="",
  1401. ):
  1402. # background task returns empty update immediately
  1403. assert update == StateUpdate()
  1404. # wait for the coroutine to start
  1405. await asyncio.sleep(0.5 if CI else 0.1)
  1406. assert len(mock_app.background_tasks) == 1
  1407. # Process another normal event
  1408. async for update in rx.app.process( # type: ignore
  1409. mock_app,
  1410. Event(
  1411. token=token,
  1412. name=f"{BackgroundTaskState.get_name()}.other",
  1413. router_data=router_data,
  1414. payload={},
  1415. ),
  1416. sid="",
  1417. headers={},
  1418. client_ip="",
  1419. ):
  1420. # other task returns delta
  1421. assert update == StateUpdate(
  1422. delta={
  1423. BackgroundTaskState.get_name(): {
  1424. "order": [
  1425. "background_task:start",
  1426. "other",
  1427. ],
  1428. }
  1429. }
  1430. )
  1431. # Explicit wait for background tasks
  1432. for task in tuple(mock_app.background_tasks):
  1433. await task
  1434. assert not mock_app.background_tasks
  1435. assert (await mock_app.state_manager.get_state(token)).order == [
  1436. "background_task:start",
  1437. "other",
  1438. "background_task:stop",
  1439. "other",
  1440. "private",
  1441. ]
  1442. @pytest.mark.asyncio
  1443. async def test_background_task_reset(mock_app: rx.App, token: str):
  1444. """Test that a background task calling reset is protected by the state proxy.
  1445. Args:
  1446. mock_app: An app that will be returned by `get_app()`
  1447. token: A token.
  1448. """
  1449. router_data = {"query": {}}
  1450. mock_app.state_manager.state = mock_app.state = BackgroundTaskState
  1451. async for update in rx.app.process( # type: ignore
  1452. mock_app,
  1453. Event(
  1454. token=token,
  1455. name=f"{BackgroundTaskState.get_name()}.background_task_reset",
  1456. router_data=router_data,
  1457. payload={},
  1458. ),
  1459. sid="",
  1460. headers={},
  1461. client_ip="",
  1462. ):
  1463. # background task returns empty update immediately
  1464. assert update == StateUpdate()
  1465. # Explicit wait for background tasks
  1466. for task in tuple(mock_app.background_tasks):
  1467. await task
  1468. assert not mock_app.background_tasks
  1469. assert (await mock_app.state_manager.get_state(token)).order == [
  1470. "reset",
  1471. ]
  1472. @pytest.mark.asyncio
  1473. async def test_background_task_no_chain():
  1474. """Test that a background task cannot be chained."""
  1475. bts = BackgroundTaskState()
  1476. with pytest.raises(RuntimeError):
  1477. await bts.bad_chain1()
  1478. with pytest.raises(RuntimeError):
  1479. await bts.bad_chain2()
  1480. def test_mutable_list(mutable_state):
  1481. """Test that mutable lists are tracked correctly.
  1482. Args:
  1483. mutable_state: A test state.
  1484. """
  1485. assert not mutable_state.dirty_vars
  1486. def assert_array_dirty():
  1487. assert mutable_state.dirty_vars == {"array"}
  1488. mutable_state._clean()
  1489. assert not mutable_state.dirty_vars
  1490. # Test all list operations
  1491. mutable_state.array.append(42)
  1492. assert_array_dirty()
  1493. mutable_state.array.extend([1, 2, 3])
  1494. assert_array_dirty()
  1495. mutable_state.array.insert(0, 0)
  1496. assert_array_dirty()
  1497. mutable_state.array.pop()
  1498. assert_array_dirty()
  1499. mutable_state.array.remove(42)
  1500. assert_array_dirty()
  1501. mutable_state.array.clear()
  1502. assert_array_dirty()
  1503. mutable_state.array += [1, 2, 3]
  1504. assert_array_dirty()
  1505. mutable_state.array.reverse()
  1506. assert_array_dirty()
  1507. mutable_state.array.sort()
  1508. assert_array_dirty()
  1509. mutable_state.array[0] = 666
  1510. assert_array_dirty()
  1511. del mutable_state.array[0]
  1512. assert_array_dirty()
  1513. # Test nested list operations
  1514. mutable_state.array[0] = [1, 2, 3]
  1515. assert_array_dirty()
  1516. mutable_state.array[0].append(4)
  1517. assert_array_dirty()
  1518. assert isinstance(mutable_state.array[0], MutableProxy)
  1519. def test_mutable_dict(mutable_state):
  1520. """Test that mutable dicts are tracked correctly.
  1521. Args:
  1522. mutable_state: A test state.
  1523. """
  1524. assert not mutable_state.dirty_vars
  1525. def assert_hashmap_dirty():
  1526. assert mutable_state.dirty_vars == {"hashmap"}
  1527. mutable_state._clean()
  1528. assert not mutable_state.dirty_vars
  1529. # Test all dict operations
  1530. mutable_state.hashmap.update({"new_key": 43})
  1531. assert_hashmap_dirty()
  1532. mutable_state.hashmap.setdefault("another_key", 66)
  1533. assert_hashmap_dirty()
  1534. mutable_state.hashmap.pop("new_key")
  1535. assert_hashmap_dirty()
  1536. mutable_state.hashmap.popitem()
  1537. assert_hashmap_dirty()
  1538. mutable_state.hashmap.clear()
  1539. assert_hashmap_dirty()
  1540. mutable_state.hashmap["new_key"] = 42
  1541. assert_hashmap_dirty()
  1542. del mutable_state.hashmap["new_key"]
  1543. assert_hashmap_dirty()
  1544. if sys.version_info >= (3, 9):
  1545. mutable_state.hashmap |= {"new_key": 44}
  1546. assert_hashmap_dirty()
  1547. # Test nested dict operations
  1548. mutable_state.hashmap["array"] = []
  1549. assert_hashmap_dirty()
  1550. mutable_state.hashmap["array"].append(1)
  1551. assert_hashmap_dirty()
  1552. mutable_state.hashmap["dict"] = {}
  1553. assert_hashmap_dirty()
  1554. mutable_state.hashmap["dict"]["key"] = 42
  1555. assert_hashmap_dirty()
  1556. mutable_state.hashmap["dict"]["dict"] = {}
  1557. assert_hashmap_dirty()
  1558. mutable_state.hashmap["dict"]["dict"]["key"] = 43
  1559. assert_hashmap_dirty()
  1560. def test_mutable_set(mutable_state):
  1561. """Test that mutable sets are tracked correctly.
  1562. Args:
  1563. mutable_state: A test state.
  1564. """
  1565. assert not mutable_state.dirty_vars
  1566. def assert_set_dirty():
  1567. assert mutable_state.dirty_vars == {"test_set"}
  1568. mutable_state._clean()
  1569. assert not mutable_state.dirty_vars
  1570. # Test all set operations
  1571. mutable_state.test_set.add(42)
  1572. assert_set_dirty()
  1573. mutable_state.test_set.update([1, 2, 3])
  1574. assert_set_dirty()
  1575. mutable_state.test_set.remove(42)
  1576. assert_set_dirty()
  1577. mutable_state.test_set.discard(3)
  1578. assert_set_dirty()
  1579. mutable_state.test_set.pop()
  1580. assert_set_dirty()
  1581. mutable_state.test_set.intersection_update([1, 2, 3])
  1582. assert_set_dirty()
  1583. mutable_state.test_set.difference_update([99])
  1584. assert_set_dirty()
  1585. mutable_state.test_set.symmetric_difference_update([102, 99])
  1586. assert_set_dirty()
  1587. mutable_state.test_set |= {1, 2, 3}
  1588. assert_set_dirty()
  1589. mutable_state.test_set &= {2, 3, 4}
  1590. assert_set_dirty()
  1591. mutable_state.test_set -= {2}
  1592. assert_set_dirty()
  1593. mutable_state.test_set ^= {42}
  1594. assert_set_dirty()
  1595. mutable_state.test_set.clear()
  1596. assert_set_dirty()
  1597. def test_mutable_custom(mutable_state):
  1598. """Test that mutable custom types derived from Base are tracked correctly.
  1599. Args:
  1600. mutable_state: A test state.
  1601. """
  1602. assert not mutable_state.dirty_vars
  1603. def assert_custom_dirty():
  1604. assert mutable_state.dirty_vars == {"custom"}
  1605. mutable_state._clean()
  1606. assert not mutable_state.dirty_vars
  1607. mutable_state.custom.foo = "bar"
  1608. assert_custom_dirty()
  1609. mutable_state.custom.array.append(42)
  1610. assert_custom_dirty()
  1611. mutable_state.custom.hashmap["key"] = 68
  1612. assert_custom_dirty()
  1613. mutable_state.custom.test_set.add(42)
  1614. assert_custom_dirty()
  1615. mutable_state.custom.custom.bar = "baz"
  1616. assert_custom_dirty()
  1617. def test_mutable_backend(mutable_state):
  1618. """Test that mutable backend vars are tracked correctly.
  1619. Args:
  1620. mutable_state: A test state.
  1621. """
  1622. assert not mutable_state.dirty_vars
  1623. def assert_custom_dirty():
  1624. assert mutable_state.dirty_vars == {"_be_custom"}
  1625. mutable_state._clean()
  1626. assert not mutable_state.dirty_vars
  1627. mutable_state._be_custom.foo = "bar"
  1628. assert_custom_dirty()
  1629. mutable_state._be_custom.array.append(42)
  1630. assert_custom_dirty()
  1631. mutable_state._be_custom.hashmap["key"] = 68
  1632. assert_custom_dirty()
  1633. mutable_state._be_custom.test_set.add(42)
  1634. assert_custom_dirty()
  1635. mutable_state._be_custom.custom.bar = "baz"
  1636. assert_custom_dirty()
  1637. @pytest.mark.parametrize(
  1638. ("copy_func",),
  1639. [
  1640. (copy.copy,),
  1641. (copy.deepcopy,),
  1642. ],
  1643. )
  1644. def test_mutable_copy(mutable_state, copy_func):
  1645. """Test that mutable types are copied correctly.
  1646. Args:
  1647. mutable_state: A test state.
  1648. copy_func: A copy function.
  1649. """
  1650. ms_copy = copy_func(mutable_state)
  1651. assert ms_copy is not mutable_state
  1652. for attr in ("array", "hashmap", "test_set", "custom"):
  1653. assert getattr(ms_copy, attr) == getattr(mutable_state, attr)
  1654. assert getattr(ms_copy, attr) is not getattr(mutable_state, attr)
  1655. ms_copy.custom.array.append(42)
  1656. assert "custom" in ms_copy.dirty_vars
  1657. if copy_func is copy.copy:
  1658. assert "custom" in mutable_state.dirty_vars
  1659. else:
  1660. assert not mutable_state.dirty_vars
  1661. @pytest.mark.parametrize(
  1662. ("copy_func",),
  1663. [
  1664. (copy.copy,),
  1665. (copy.deepcopy,),
  1666. ],
  1667. )
  1668. def test_mutable_copy_vars(mutable_state, copy_func):
  1669. """Test that mutable types are copied correctly.
  1670. Args:
  1671. mutable_state: A test state.
  1672. copy_func: A copy function.
  1673. """
  1674. for attr in ("array", "hashmap", "test_set", "custom"):
  1675. var_orig = getattr(mutable_state, attr)
  1676. var_copy = copy_func(var_orig)
  1677. assert var_orig is not var_copy
  1678. assert var_orig == var_copy
  1679. # copied vars should never be proxies, as they by definition are no longer attached to the state.
  1680. assert not isinstance(var_copy, MutableProxy)
  1681. def test_duplicate_substate_class(duplicate_substate):
  1682. with pytest.raises(ValueError):
  1683. duplicate_substate()