From dd0c5fe28378a2b3aafd9619bfa72070227ad05a Mon Sep 17 00:00:00 2001 From: ViditOstwal Date: Fri, 18 Sep 2026 15:36:46 +0530 Subject: [PATCH] test(flows): cover MongoDB persistence isolation --- .../tests/test_flow_persistence_mongodb.py | 113 +++++++++++++++++- 1 file changed, 111 insertions(+), 2 deletions(-) diff --git a/lib/crewai/tests/test_flow_persistence_mongodb.py b/lib/crewai/tests/test_flow_persistence_mongodb.py index daaa50b47..c1c1393c5 100644 --- a/lib/crewai/tests/test_flow_persistence_mongodb.py +++ b/lib/crewai/tests/test_flow_persistence_mongodb.py @@ -143,7 +143,9 @@ class _FakeClient: return _FakeSession(self) -def _patch_client(monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: +def _patch_client( + monkeypatch: pytest.MonkeyPatch, *, reuse_client: bool = False +) -> dict[str, Any]: """Patch ``pymongo.MongoClient`` to build in-memory fakes. Returns a dict capturing the connection string and the created fake client @@ -152,9 +154,16 @@ def _patch_client(monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: import pymongo created: dict[str, Any] = {} + shared_client: _FakeClient | None = None def factory(conn: str, *args: Any, **kwargs: Any) -> _FakeClient: - client = _FakeClient(conn) + nonlocal shared_client + if reuse_client: + if shared_client is None: + shared_client = _FakeClient(conn) + client = shared_client + else: + client = _FakeClient(conn) created["conn"] = conn created["client"] = client return client @@ -198,6 +207,38 @@ def test_save_state_tags_incrementing_seq(monkeypatch: pytest.MonkeyPatch) -> No assert last_sort == [("seq", -1)] +def test_persistences_sharing_a_state_collection_share_its_counter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Separate persistence instances keep one ordering sequence per collection.""" + created = _patch_client(monkeypatch, reuse_client=True) + first = MongoDbFlowPersistence(CONN) + second = MongoDbFlowPersistence(CONN) + + first.save_state("flow-1", "first", {"counter": 1}) + second.save_state("flow-1", "second", {"counter": 2}) + + states = created["client"]._db["flow_states"] + assert [doc["seq"] for doc in states.docs] == [1, 2] + assert first.load_state("flow-1") == {"counter": 2} + + +def test_different_state_collections_use_independent_counters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Each configured state collection has an independent ordering sequence.""" + created = _patch_client(monkeypatch, reuse_client=True) + default_states = MongoDbFlowPersistence(CONN) + billing_states = MongoDbFlowPersistence(CONN, states_collection="billing_states") + + default_states.save_state("flow-1", "default", {"counter": 1}) + billing_states.save_state("flow-1", "billing", {"counter": 1}) + + database = created["client"]._db + assert database["flow_states"].docs[0]["seq"] == 1 + assert database["billing_states"].docs[0]["seq"] == 1 + + def test_save_state_assigns_sequence_and_inserts_in_one_transaction( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -333,6 +374,38 @@ def test_pending_feedback_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None: assert created["client"].transactions_started == 1 +def test_pending_feedback_replaces_existing_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Saving feedback twice matches SQLite's INSERT OR REPLACE behavior.""" + _patch_client(monkeypatch) + persistence = MongoDbFlowPersistence(CONN) + first_context = PendingFeedbackContext( + flow_id="flow-1", + flow_class="MyFlow", + method_name="first_review", + method_output="first draft", + message="First question?", + ) + second_context = PendingFeedbackContext( + flow_id="flow-1", + flow_class="MyFlow", + method_name="second_review", + method_output="second draft", + message="Second question?", + ) + + persistence.save_pending_feedback("flow-1", first_context, {"counter": 1}) + persistence.save_pending_feedback("flow-1", second_context, {"counter": 2}) + + loaded = persistence.load_pending_feedback("flow-1") + assert loaded is not None + state, context = loaded + assert state == {"counter": 2} + assert context.method_name == "second_review" + assert context.message == "Second question?" + + def test_pending_feedback_write_failure_rolls_back_state( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -386,6 +459,42 @@ def test_persisted_flow_restores_latest_mongodb_state( assert persistence.load_state(flow_id) == {"id": flow_id, "counter": 4} +def test_mongodb_fork_restores_source_state_without_mutating_source( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Mongo persistence follows SQLite's restore_from_state_id fork semantics.""" + _patch_client(monkeypatch) + persistence = MongoDbFlowPersistence(CONN) + + class State(FlowState): + counter: int = 0 + + class ForkableFlow(Flow[State]): + @start() + @persist(persistence) + def step(self) -> None: + self.state.counter += 1 + + source = ForkableFlow(persistence=persistence) + source.kickoff() + source_id = source.state.id + + resumed_source = ForkableFlow(persistence=persistence) + resumed_source.kickoff(inputs={"id": source_id}) + assert persistence.load_state(source_id) == {"id": source_id, "counter": 2} + + fork = ForkableFlow(persistence=persistence) + fork.kickoff(restore_from_state_id=source_id) + + assert fork.state.id != source_id + assert fork.state.counter == 3 + assert persistence.load_state(source_id) == {"id": source_id, "counter": 2} + assert persistence.load_state(fork.state.id) == { + "id": fork.state.id, + "counter": 3, + } + + def test_mongodb_from_pending_resumes_flow(monkeypatch: pytest.MonkeyPatch) -> None: _patch_client(monkeypatch) persistence = MongoDbFlowPersistence(CONN)