From b969ab3ff7401ec13b56629ce2ffd05368fa47a0 Mon Sep 17 00:00:00 2001 From: ponyisi Date: Fri, 6 Dec 2024 10:45:31 -0600 Subject: [PATCH] Avoid race condition between status monitoring and download tasks (#532) Add condition to exit polling loop; update tests --- servicex/query_core.py | 3 ++- tests/test_dataset.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/servicex/query_core.py b/servicex/query_core.py index 6a21d99f..d58841c3 100644 --- a/servicex/query_core.py +++ b/servicex/query_core.py @@ -562,7 +562,8 @@ async def get_signed_url( # signing urls for a previous transform then we know it is complete as well if cached_record or ( self.current_status - and self.current_status.status in DONE_STATUS + and (self.current_status.status in DONE_STATUS + and self.current_status.files_completed == len(files_seen)) ): break diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 742fe059..31e3a098 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -116,7 +116,7 @@ async def test_download_files(python_dataset): progress_mock = Mock() python_dataset.minio_polling_interval = 0 python_dataset.minio = minio_mock - python_dataset.current_status = Mock(status="Complete") + python_dataset.current_status = Mock(status="Complete", files_completed=2) python_dataset.configuration.shortened_downloaded_filename = False result_uris = await python_dataset.download_files( @@ -142,7 +142,7 @@ async def test_download_files_with_signed_urls(python_dataset): python_dataset.minio_polling_interval = 0 python_dataset.minio = minio_mock - python_dataset.current_status = Mock(status="Complete") + python_dataset.current_status = Mock(status="Complete", files_completed=2) python_dataset.configuration.shortened_downloaded_filename = False result_uris = await python_dataset.download_files(