added tests and cache for ndjson file fetching

Author: Matt Sokoloff

labelbox/client.py

@@ -253,13 +253,15 @@ def upload_file(self, path: str) -> str:
     def upload_data(self,
                     content: bytes,
                     filename: str = None,
-                    content_type: str = None) -> str:
+                    content_type: str = None,
+                    sign: bool = False) -> str:
         """ Uploads the given data (bytes) to Labelbox.
 
         Args:
             content: bytestring to upload
             filename: name of the upload
             content_type: content type of data uploaded
+            sign: whether or not to sign the url
 
         Returns:
             str, the URL of uploaded data.
@@ -274,7 +276,7 @@ def upload_data(self,
                     "variables": {
                         "file": None,
                         "contentLength": len(content),
-                        "sign": False
+                        "sign": sign
                     },
                     "query":
                         """mutation UploadFile($file: Upload!, $contentLength: Int!,

labelbox/schema/bulk_import_request.py

@@ -1,6 +1,7 @@
 import json
 import time
 from uuid import UUID, uuid4
+import functools
 
 import logging
 from pathlib import Path
@@ -114,35 +115,58 @@ class BulkImportRequest(DbObject):
     created_by = Relationship.ToOne("User", False, "created_by")
 
     @property
-    def inputs(self):
+    def inputs(self) -> Optional[List[Dict[str, str]]]:
         """
         Inputs for each individual annotation uploaded.
-        * This should match the ndjson annotations that you have uploaded. 
-        * This information will expire after 24 hours.
-        """        
-        return  self._fetch_remote_ndjson(self.input_file_url)
+        This should match the ndjson annotations that you have uploaded. 
         
+        Returns:
+            Uploaded ndjsons.
+
+        * This information will expire after 24 hours.    
+        """
+        return self._fetch_remote_ndjson(self.input_file_url)
+
     @property
-    def errors(self):
+    def errors(self) -> Optional[List[Dict[str, str]]]:
         """
         Errors for each individual annotation uploaded.
-        * Returns an empty list if there are no errors and None if the update is still running.
+
+        Returns:
+            Empty list if there are no errors and None if the update is still running.
+            If there are errors, and the job has completed then a list of dicts containing the error messages will be returned.
+
         * This information will expire after 24 hours.        
         """
-        return  self._fetch_remote_ndjson(self.error_file_url)
+        return self._fetch_remote_ndjson(self.error_file_url)
 
     @property
-    def statuses(self):
+    def statuses(self) -> Optional[List[Dict[str, str]]]:
         """
         Status for each individual annotation uploaded.
-        * Returns a status for each row if the upload is done running and was successful. Otherwise it returns None.
+
+        Returns:
+            A status for each annotation if the upload is done running and was successful. Otherwise it returns None.
+
         * This information will expire after 24 hours.        
         """
-        return  self._fetch_remote_ndjson(self.status_file_url)
+        return self._fetch_remote_ndjson(self.status_file_url)
+
+    @functools.lru_cache()
+    def _fetch_remote_ndjson(
+            self, url: Optional[str]) -> Optional[List[Dict[str, str]]]:
+        """
+        Fetches the remote ndjson file and caches the results.
 
-    def _fetch_remote_ndjson(self, url):
+        Args:
+            url (str): either the input_file_url, error_file_url, status_file_url, or None
+                urls are None when the file is unavailable.
+        Returns:
+            None if the url is None or the ndjson as a list of dicts.
+        """
         if url is not None:
             return ndjson.loads(requests.get(url).text)
+        return None
 
     def refresh(self) -> None:
         """Synchronizes values of all fields with the database.

tests/integration/bulk_import/test_bulk_import_request.py

@@ -121,20 +121,26 @@ def test_validate_ndjson_uuid(tmp_path, configured_project, predictions):
 
 
 @pytest.mark.slow
-def test_wait_till_done(configured_project):
+def test_wait_till_done(rectangle_inference, configured_project):
     name = str(uuid.uuid4())
-    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
+    url = configured_project.client.upload_data(content=ndjson.dumps(
+        [rectangle_inference]),
+                                                sign=True)
     bulk_import_request = configured_project.upload_annotations(name=name,
                                                                 annotations=url,
                                                                 validate=False)
 
     bulk_import_request.wait_until_done()
-
-    assert (bulk_import_request.state == BulkImportRequestState.FINISHED or
-            bulk_import_request.state == BulkImportRequestState.FAILED)
-
-    assert bulk_import_request.errors is not None
-    assert bulk_import_request.inputs is not None
+    assert bulk_import_request.state == BulkImportRequestState.FINISHED
+
+    #Check that the status files are being returned as expected
+    assert len(bulk_import_request.errors) == 0
+    assert len(bulk_import_request.inputs) == 1
+    assert bulk_import_request.inputs[0]['uuid'] == rectangle_inference['uuid']
+    assert len(bulk_import_request.statuses) == 1
+    assert bulk_import_request.statuses[0]['status'] == 'SUCCESS'
+    assert bulk_import_request.statuses[0]['uuid'] == rectangle_inference[
+        'uuid']
 
 
 def assert_file_content(url: str, predictions):