diff --git a/burr/core/application.py b/burr/core/application.py index 91ef8b1b6..cb19fe95f 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -2707,8 +2707,9 @@ def _init_state_from_persister( self.start = last_position self.reset_to_entrypoint() else: - # self.start is already set to the default. We don't need to do anything. - pass + # self.start is already set to the default, but the loaded state still has the + # prior step, which would otherwise decide the next action. Clear it. + self.reset_to_entrypoint() def _load_from_sync_persister(self): """Loads from the set sync persister and into this current object. diff --git a/tests/core/test_application.py b/tests/core/test_application.py index efb7d205b..8f2d7c81d 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -67,6 +67,7 @@ BaseStateLoader, BaseStatePersister, DevNullPersister, + InMemoryPersister, PersistedStateData, PersisterHookAsync, SQLLitePersister, @@ -4454,3 +4455,76 @@ def noop(state: State) -> State: app = builder.build() assert app.state["x"] == 100 + + +@pytest.mark.parametrize("last_step_failed", [False, True]) +@pytest.mark.parametrize( + "resume_at_next_action,expected_next_action", + [(True, "result"), (False, "counter")], +) +def test_initialize_from_resume_at_next_action( + resume_at_next_action, expected_next_action, last_step_failed +): + """resume_at_next_action=False should start at the default entrypoint, keeping the loaded state, + whether the last saved step completed or failed""" + persister = InMemoryPersister() + result = base_broken_action if last_step_failed else Result("count") + + def build(**initialize_kwargs): + return ( + ApplicationBuilder() + .with_actions(counter=base_counter_action, result=result) + .with_transitions(("counter", "result", default)) + .initialize_from( + persister, default_state={}, default_entrypoint="counter", **initialize_kwargs + ) + .with_identifiers(app_id="app", partition_key="pk") + .with_state_persister(persister) + .build() + ) + + build(resume_at_next_action=True).step() + if last_step_failed: + with pytest.raises(BrokenStepException): + build(resume_at_next_action=True).step() + app = build(resume_at_next_action=resume_at_next_action) + assert app.get_next_action().name == expected_next_action + assert app.state["count"] == 1 + if not resume_at_next_action: + assert PRIOR_STEP not in app.state.get_all() + + +@pytest.mark.parametrize("last_step_failed", [False, True]) +@pytest.mark.parametrize( + "resume_at_next_action,expected_next_action", + [(True, "result"), (False, "counter")], +) +async def test_initialize_from_resume_at_next_action_async( + resume_at_next_action, expected_next_action, last_step_failed +): + """Same as above, for the async builder""" + persister = AsyncInMemoryPersister() + result = base_broken_action_async if last_step_failed else Result("count") + + async def abuild(**initialize_kwargs): + return await ( + ApplicationBuilder() + .with_actions(counter=base_counter_action_async, result=result) + .with_transitions(("counter", "result", default)) + .initialize_from( + persister, default_state={}, default_entrypoint="counter", **initialize_kwargs + ) + .with_identifiers(app_id="app", partition_key="pk") + .with_state_persister(persister) + .abuild() + ) + + await (await abuild(resume_at_next_action=True)).astep() + if last_step_failed: + with pytest.raises(BrokenStepException): + await (await abuild(resume_at_next_action=True)).astep() + app = await abuild(resume_at_next_action=resume_at_next_action) + assert app.get_next_action().name == expected_next_action + assert app.state["count"] == 1 + if not resume_at_next_action: + assert PRIOR_STEP not in app.state.get_all()