Avoid race condition between status monitoring and download tasks (#532)

Add condition to exit polling loop; update tests
ssl-hep · Dec 6, 2024 · b969ab3 · b969ab3
1 parent 2e2af1b
commit b969ab3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git 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
@@ -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(