diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index 807f1e75..3e8c790a 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/taskiq/brokers/inmemory_broker.py @@ -44,7 +44,8 @@ async def set_result(self, task_id: str, result: TaskiqResult[_ReturnType]) -> N :param result: result of an execution. """ if ( - self.max_stored_results != -1 + task_id not in self.results + and self.max_stored_results != -1 and len(self.results) >= self.max_stored_results ): self.results.popitem(last=False) @@ -95,7 +96,8 @@ async def set_progress( :param progress: task execution progress """ if ( - self.max_stored_results != -1 + task_id not in self.progress + and self.max_stored_results != -1 and len(self.progress) >= self.max_stored_results ): self.progress.popitem(last=False) diff --git a/tests/brokers/test_inmemory_result_backend.py b/tests/brokers/test_inmemory_result_backend.py new file mode 100644 index 00000000..c2a4641e --- /dev/null +++ b/tests/brokers/test_inmemory_result_backend.py @@ -0,0 +1,74 @@ +import pytest + +from taskiq.brokers.inmemory_broker import InmemoryResultBackend +from taskiq.depends.progress_tracker import TaskProgress +from taskiq.result import TaskiqResult + + +@pytest.mark.parametrize("capacity", [1, 2, -1]) +@pytest.mark.parametrize("update_first", [True, False]) +async def test_result_updates_preserve_cached_tasks( + capacity: int, + *, + update_first: bool, +) -> None: + backend = InmemoryResultBackend[int](max_stored_results=capacity) + task_ids = ["first"] if capacity == 1 else ["first", "second"] + for task_id in task_ids: + await backend.set_result( + task_id, + TaskiqResult(is_err=False, return_value=0, execution_time=0), + ) + + updated_id = task_ids[0] if update_first else task_ids[-1] + for value in (1, 2): + await backend.set_result( + updated_id, + TaskiqResult(is_err=False, return_value=value, execution_time=0), + ) + for task_id in task_ids: + assert await backend.is_result_ready(task_id) + assert (await backend.get_result(updated_id)).return_value == value + + await backend.set_result( + "new", + TaskiqResult(is_err=False, return_value=3, execution_time=0), + ) + assert await backend.is_result_ready("new") + assert await backend.is_result_ready("first") is (capacity == -1) + if capacity != 1: + assert await backend.is_result_ready("second") + + +@pytest.mark.parametrize("capacity", [1, 2, -1]) +@pytest.mark.parametrize("update_first", [True, False]) +async def test_progress_updates_preserve_cached_tasks( + capacity: int, + *, + update_first: bool, +) -> None: + backend = InmemoryResultBackend[int](max_stored_results=capacity) + task_ids = ["first"] if capacity == 1 else ["first", "second"] + for task_id in task_ids: + await backend.set_progress( + task_id, + TaskProgress(state="STARTED", meta=0), + ) + + updated_id = task_ids[0] if update_first else task_ids[-1] + for value in (1, 2): + await backend.set_progress( + updated_id, + TaskProgress(state="STARTED", meta=value), + ) + for task_id in task_ids: + assert await backend.get_progress(task_id) is not None + progress = await backend.get_progress(updated_id) + assert progress is not None + assert progress.meta == value + + await backend.set_progress("new", TaskProgress(state="STARTED", meta=3)) + assert await backend.get_progress("new") is not None + assert (await backend.get_progress("first") is not None) is (capacity == -1) + if capacity != 1: + assert await backend.get_progress("second") is not None