From b75ef482bf660499c5057a910a45693f3d6f77af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Moura?= Date: Sun, 19 Jan 2025 00:46:26 -0800 Subject: [PATCH] refactoring test --- tests/test_flow_default_override.py | 44 +++++++++++++++++------------ 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/tests/test_flow_default_override.py b/tests/test_flow_default_override.py index 28053c249..b2bcfa9c5 100644 --- a/tests/test_flow_default_override.py +++ b/tests/test_flow_default_override.py @@ -14,42 +14,50 @@ from crewai.flow.persistence.sqlite import SQLiteFlowPersistence class PoemState(FlowState): """Test state model with default values that should be overridden.""" sentence_count: int = 1000 # Default that should be overridden - poem: str = "" has_set_count: bool = False # Track whether we've set the count -def test_default_value_override(tmp_path): +def test_default_value_override(): """Test that persisted state values override class defaults.""" - db_path = os.path.join(tmp_path, "test_flows.db") - persistence = SQLiteFlowPersistence(db_path) - - @persist(persistence) + + @persist() class PoemFlow(Flow[PoemState]): initial_state = PoemState - + @start() def set_sentence_count(self): + print("Setting sentence count") + print(self.state) # Only set sentence_count on first run, not when loading from persistence - if not self.state.has_set_count: + if self.state.has_set_count and self.state.sentence_count == 2: + self.state.sentence_count = 3 + elif self.state.has_set_count and self.state.sentence_count == 1000: + self.state.sentence_count = 1000 + elif self.state.has_set_count and self.state.sentence_count == 5: + self.state.sentence_count = 5 + else: self.state.sentence_count = 2 self.state.has_set_count = True - + # First run - should set sentence_count to 2 - # First run - should set sentence_count to 2 - flow1 = PoemFlow(persistence=persistence) + flow1 = PoemFlow() flow1.kickoff() original_uuid = flow1.state.id assert flow1.state.sentence_count == 2 - + # Second run - should load sentence_count=2 instead of default 1000 - flow2 = PoemFlow(persistence=persistence) + flow2 = PoemFlow() flow2.kickoff(inputs={"id": original_uuid}) - assert flow2.state.sentence_count == 2 # Should load 2, not default 1000 - + assert flow2.state.sentence_count == 3 # Should load 2, not default 1000 + + # Third run - should not load sentence_count=2 instead of default 1000 + flow2 = PoemFlow() + flow2.kickoff(inputs={"has_set_count": True}) + assert flow2.state.sentence_count == 1000 # Should load 1000, not 2 + # Third run - explicit override should work flow3 = PoemFlow( - persistence=persistence, id=original_uuid, - sentence_count=3, # Override persisted value + sentence_count=5, # Override persisted value ) - assert flow3.state.sentence_count == 3 # Should use override value + assert flow3.state.sentence_count == 5 # Should use override value