resolve breaking tests

Author: MattShirley

servicex/query_core.py

@@ -319,7 +319,6 @@ def transform_complete(task: Task):
             else None
         )
 
-        begin_at = datetime.datetime.now(tz=datetime.timezone.utc)
         if not cached_record:
 
             if self.cache.is_transform_request_submitted(sx_request_hash):
@@ -348,7 +347,6 @@ def transform_complete(task: Task):
                 expandable_progress,
                 download_progress,
                 cached_record,
-                begin_at,
             )
         )
 
@@ -524,7 +522,6 @@ async def download_files(
         progress: ExpandableProgress,
         download_progress: TaskID,
         cached_record: Optional[TransformedResults],
-        begin_at: datetime.datetime,
     ) -> List[str]:
         """
         Task to monitor the list of files in the transform output's bucket. Any new files
@@ -560,17 +557,16 @@ async def get_signed_url(
             if progress:
                 progress.advance(task_id=download_progress, task_type="Download")
 
+        later_than: datetime.datetime | None = None
         while True:
             if not cached_record:
                 await asyncio.sleep(self.minio_polling_interval)
             if self.minio:
                 # if self.minio exists, self.current_status will too
                 if self.current_status.files_completed > len(files_seen):
-                    new_begin_at = datetime.datetime.now(tz=datetime.timezone.utc)
                     files = await self.servicex.get_transformation_results(
-                        self.current_status.request_id, begin_at
+                        self.current_status.request_id, later_than
                     )
-                    begin_at = new_begin_at
 
                     for file in files:
                         filename = file.filename
@@ -601,6 +597,9 @@ async def get_signed_url(
                                 )  # NOQA 501
                             files_seen.add(filename)
 
+                            if later_than is None or file.created_at > later_than:
+                                later_than = file.created_at
+
             # Once the transform is complete and all files are seen we can stop polling.
             # Also, if we are just downloading or signing urls for a previous transform
             # then we know it is complete as well

servicex/servicex_adapter.py

@@ -52,6 +52,7 @@ class AuthorizationError(BaseException):
 
 @dataclass
 class ServiceXFile:
+    created_at: datetime.datetime
     filename: str
 
 
@@ -236,13 +237,13 @@ async def delete_transform(self, transform_id=None):
                     )
 
     async def get_transformation_results(
-        self, request_id: str, begin_at: datetime.datetime
+        self, request_id: str, later_than: datetime.datetime | None = None
     ):
         headers = await self._get_authorization()
         url = self.url + f"/servicex/transformation/{request_id}/results"
-        params = {
-            "begin_at": begin_at.isoformat(),
-        }
+        params = {}
+        if later_than:
+            params["later_than"] = later_than.isoformat()
 
         async with ClientSession() as session:
             async with session.get(headers=headers, url=url, params=params) as r:
@@ -261,7 +262,10 @@ async def get_transformation_results(
                 data = await r.json()
                 response = list()
                 for result in data.get("results", []):
-                    file = ServiceXFile(filename=result["file-path"].replace("/", ":"))
+                    file = ServiceXFile(
+                        filename=result["file-path"].replace("/", ":"),
+                        created_at=datetime.datetime.fromisoformat(result["created_at"]),
+                    )
                     response.append(file)
                 return response
 

tests/test_dataset.py

@@ -45,6 +45,8 @@
 )
 from rich.progress import Progress
 
+from servicex.servicex_adapter import ServiceXFile
+
 
 @pytest.mark.asyncio
 async def test_as_signed_urls_happy(transformed_result):
@@ -121,16 +123,15 @@ async def test_as_files_cached(transformed_result, python_dataset):
 @pytest.mark.asyncio
 async def test_download_files(python_dataset):
     signed_urls_only = False
-    begin_at = datetime.datetime.now(tz=datetime.timezone.utc)
     download_progress = "download_task_id"
     minio_mock = AsyncMock()
     config = Configuration(cache_path="temp_dir", api_endpoints=[])
     python_dataset.configuration = config
     python_dataset.servicex = AsyncMock()
     python_dataset.servicex.get_transformation_results = AsyncMock()
     python_dataset.servicex.get_transformation_results.return_value = [
-        Mock(filename="file1.txt"),
-        Mock(filename="file2.txt"),
+        ServiceXFile(filename="file1.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ServiceXFile(filename="file2.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
     ]
 
     minio_mock.download_file.return_value = Path("/path/to/downloaded_file")
@@ -143,7 +144,7 @@ async def test_download_files(python_dataset):
     python_dataset.configuration.shortened_downloaded_filename = False
 
     result_uris = await python_dataset.download_files(
-        signed_urls_only, progress_mock, download_progress, None, begin_at
+        signed_urls_only, progress_mock, download_progress, None
     )
     minio_mock.download_file.assert_awaited()
     minio_mock.get_signed_url.assert_not_awaited()
@@ -153,7 +154,6 @@ async def test_download_files(python_dataset):
 @pytest.mark.asyncio
 async def test_download_files_with_signed_urls(python_dataset):
     signed_urls_only = True
-    begin_at = datetime.datetime.now(tz=datetime.timezone.utc)
     download_progress = "download_task_id"
     minio_mock = AsyncMock()
     config = Configuration(cache_path="temp_dir", api_endpoints=[])
@@ -165,8 +165,8 @@ async def test_download_files_with_signed_urls(python_dataset):
     python_dataset.servicex = AsyncMock()
     python_dataset.servicex.get_transformation_results = AsyncMock()
     python_dataset.servicex.get_transformation_results.return_value = [
-        Mock(filename="file1.txt"),
-        Mock(filename="file2.txt"),
+        ServiceXFile(filename="file1.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ServiceXFile(filename="file2.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
     ]
 
     python_dataset.minio_polling_interval = 0
@@ -175,7 +175,7 @@ async def test_download_files_with_signed_urls(python_dataset):
     python_dataset.configuration.shortened_downloaded_filename = False
 
     result_uris = await python_dataset.download_files(
-        signed_urls_only, progress_mock, download_progress, None, begin_at
+        signed_urls_only, progress_mock, download_progress, None
     )
     minio_mock.download_file.assert_not_called()
     minio_mock.get_signed_url.assert_called()

tests/test_servicex_adapter.py

@@ -29,6 +29,7 @@
 import tempfile
 import time
 import datetime
+from unittest import result
 from unittest.mock import patch, AsyncMock
 
 import httpx
@@ -512,8 +513,8 @@ async def test_get_transformation_results_success(get, servicex):
     get.return_value.__aenter__.return_value.json = AsyncMock(
         return_value={
             "results": [
-                {"file-path": "file1.txt"},
-                {"file-path": "file2.txt"},
+                {"file-path": "file1.txt", "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat()},
+                {"file-path": "file2.txt", "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat()},
             ]
         }
     )
@@ -528,7 +529,7 @@ async def test_get_transformation_results_success(get, servicex):
         url=f"https://servicex.org/servicex/transformation/{request_id}/results",
         headers={},
         params={
-            "begin_at": now.isoformat(),
+            "later_than": now.isoformat(),
         },
     )
 
@@ -549,7 +550,7 @@ async def test_get_transformation_results_not_found(
         url=f"https://servicex.org/servicex/transformation/{request_id}/results",
         headers={},
         params={
-            "begin_at": now.isoformat(),
+            "later_than": now.isoformat(),
         },
     )
 
@@ -570,7 +571,7 @@ async def test_get_transformation_results_not_authorized(
         url=f"https://servicex.org/servicex/transformation/{request_id}/results",
         headers={},
         params={
-            "begin_at": now.isoformat(),
+            "later_than": now.isoformat(),
         },
     )
 
@@ -591,6 +592,6 @@ async def test_get_transformation_results_server_error(
         url=f"https://servicex.org/servicex/transformation/{request_id}/results",
         headers={},
         params={
-            "begin_at": now.isoformat(),
+            "later_than": now.isoformat(),
         },
     )

tests/test_servicex_dataset.py

@@ -25,6 +25,7 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+import datetime
 import tempfile
 from typing import List
 from unittest.mock import AsyncMock, patch, Mock
@@ -46,6 +47,7 @@
 )
 from servicex.query_cache import QueryCache
 from servicex.query_core import ServiceXException, Query
+from servicex.servicex_adapter import ServiceXFile
 from servicex.servicex_client import ServiceXClient
 from servicex.uproot_raw.uproot_raw import UprootRawQuery
 
@@ -208,10 +210,15 @@ async def test_submit(mocker):
     servicex.submit_transform = AsyncMock()
     servicex.submit_transform.return_value = {"request_id": '123-456-789"'}
 
-    servicex.get_transformation_results = AsyncMock(
-        side_effect=[[file1], [file1, file2]]
-    )
-
+    servicex.get_transformation_results = AsyncMock(side_effect=[
+        [
+            ServiceXFile(filename="file1", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ],
+        [
+            ServiceXFile(filename="file1", created_at=datetime.datetime.now(datetime.timezone.utc)),
+            ServiceXFile(filename="file2", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ]
+    ])
     servicex.get_transform_status = AsyncMock()
     servicex.get_transform_status.side_effect = [
         transform_status1,
@@ -256,7 +263,14 @@ async def test_submit_partial_success(mocker):
     servicex = AsyncMock()
     servicex.submit_transform = AsyncMock()
     servicex.submit_transform.return_value = {"request_id": '123-456-789"'}
-    servicex.get_transformation_results = AsyncMock(side_effect=[[file1], [file1]])
+    servicex.get_transformation_results = AsyncMock(side_effect=[
+        [
+            ServiceXFile(filename="file1", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ],
+        [
+            ServiceXFile(filename="file1", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ]
+    ])
 
     servicex.get_transform_status = AsyncMock()
     servicex.get_transform_status.side_effect = [
@@ -304,8 +318,8 @@ async def test_use_of_cache(mocker):
     servicex.submit_transform.return_value = {"request_id": '123-456-789"'}
     servicex.get_transformation_results = AsyncMock()
     servicex.get_transformation_results.return_value = [
-        Mock(filename="file1.txt"),
-        Mock(filename="file2.txt"),
+        ServiceXFile(filename="file1.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ServiceXFile(filename="file2.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
     ]
     servicex.get_transform_status = AsyncMock()
     servicex.get_transform_status.side_effect = [
@@ -610,8 +624,8 @@ async def test_use_of_ignore_cache(mocker, servicex):
     )
     servicex.get_transformation_results = AsyncMock()
     servicex.get_transformation_results.return_value = [
-        Mock(filename="file1.txt"),
-        Mock(filename="file2.txt"),
+        ServiceXFile(filename="file1.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
+        ServiceXFile(filename="file2.txt", created_at=datetime.datetime.now(datetime.timezone.utc)),
     ]
     # Prepare Minio
     mock_minio = AsyncMock()