Merge pull request #774 from simvue-io/hotfix/fix-obj-get-count

Fix count limiting on object retrieval

Author: wk9874

simvue/api/objects/base.py

@@ -365,7 +365,7 @@ def ids(
         """
         _class_instance = cls(_read_only=True, _local=True)
         _count: int = 0
-        for response in cls._get_all_objects(offset):
+        for response in cls._get_all_objects(offset, count=count):
             if (_data := response.get("data")) is None:
                 raise RuntimeError(
                     f"Expected key 'data' for retrieval of {_class_instance.__class__.__name__.lower()}s"
@@ -404,7 +404,7 @@ def get(
         """
         _class_instance = cls(_read_only=True, _local=True)
         _count: int = 0
-        for _response in cls._get_all_objects(offset, **kwargs):
+        for _response in cls._get_all_objects(offset, count=count, **kwargs):
             if count and _count > count:
                 return
             if (_data := _response.get("data")) is None:
@@ -438,7 +438,7 @@ def count(cls, **kwargs) -> int:
 
     @classmethod
     def _get_all_objects(
-        cls, offset: int | None, **kwargs
+        cls, offset: int | None, count: int | None, **kwargs
     ) -> typing.Generator[dict, None, None]:
         _class_instance = cls(_read_only=True)
         _url = f"{_class_instance._base_url}"
@@ -448,7 +448,7 @@ def _get_all_objects(
             _label = _label[:-1]
 
         for response in get_paginated(
-            _url, headers=_class_instance._headers, offset=offset, **kwargs
+            _url, headers=_class_instance._headers, offset=offset, count=count, **kwargs
         ):
             yield get_json_from_response(
                 response=response,

simvue/api/request.py

@@ -281,6 +281,7 @@ def get_paginated(
     timeout: int = DEFAULT_API_TIMEOUT,
     json: dict[str, typing.Any] | None = None,
     offset: int | None = None,
+    count: int | None = None,
     **params,
 ) -> typing.Generator[requests.Response, None, None]:
     """Paginate results of a server query.
@@ -309,7 +310,7 @@ def get_paginated(
                 url=url,
                 headers=headers,
                 params=(params or {})
-                | {"count": MAX_ENTRIES_PER_PAGE, "start": _offset},
+                | {"count": count or MAX_ENTRIES_PER_PAGE, "start": _offset},
                 timeout=timeout,
                 json=json,
             )
@@ -320,4 +321,5 @@ def get_paginated(
         yield _response
         _offset += MAX_ENTRIES_PER_PAGE
 
-    yield _response
+        if count and _offset > count:
+            break

tests/unit/test_folder.py

@@ -38,22 +38,32 @@ def test_folder_creation_offline() -> None:
 
     with _folder._local_staging_file.open() as in_f:
         _local_data = json.load(in_f)
-        
+
     assert  _folder._local_staging_file.name.split(".")[0] == _folder.id
     assert _local_data.get("path", None) == _path
-        
+
     sender(_folder._local_staging_file.parents[1], 2, 10, ["folders"])
     time.sleep(1)
     client = Client()
-    
+
     _folder_new = client.get_folder(_path)
     assert _folder_new.path == _path
-    
+
     _folder_new.delete()
-    
+
     assert not _folder._local_staging_file.exists()
 
 
+@pytest.mark.api
+@pytest.mark.online
+def test_get_folder() -> None:
+    _uuid: str = f"{uuid.uuid4()}".split("-")[0]
+    _folder_name = f"/simvue_unit_testing/{_uuid}"
+    _folder_1 = Folder.new(path=f"{_folder_name}/dir_1")
+    _folder_2 = Folder.new(path=f"{_folder_name}/dir_2")
+    assert len(list(Folder._get_all_objects(count=2, offset=None))) == 2
+
+
 @pytest.mark.api
 @pytest.mark.online
 def test_folder_modification_online() -> None:
@@ -85,34 +95,34 @@ def test_folder_modification_offline() -> None:
     _tags = ["testing", "api"]
     _folder = Folder.new(path=_path, offline=True)
     _folder.commit()
-    
+
     sender(_folder._local_staging_file.parents[1], 2, 10, ["folders"])
     time.sleep(1)
-    
+
     client = Client()
     _folder_online = client.get_folder(_path)
     assert _folder_online.path == _path
-    
+
     _folder_new = Folder(identifier=_folder.id)
     _folder_new.read_only(False)
     _folder_new.tags = _tags
     _folder_new.description = _description
     _folder_new.commit()
-    
+
     with _folder._local_staging_file.open() as in_f:
         _local_data = json.load(in_f)
     assert  _folder._local_staging_file.name.split(".")[0] == _folder.id
     assert _local_data.get("description", None) == _description
     assert _local_data.get("tags", None) == _tags
-        
+
     sender(_folder._local_staging_file.parents[1], 2, 10, ["folders"])
     time.sleep(1)
-    
+
     _folder_online.refresh()
     assert _folder_online.path == _path
     assert _folder_online.description == _description
     assert _folder_online.tags == _tags
-    
+
     _folder_online.read_only(False)
     _folder_online.delete()
 

tests/unit/test_run.py

@@ -39,17 +39,17 @@ def test_run_creation_offline() -> None:
         _local_data = json.load(in_f)
     assert _local_data.get("name") == f"simvue_offline_run_{_uuid}"
     assert _local_data.get("folder") == _folder_name
-    
+
     sender(_run._local_staging_file.parents[1], 1, 10, ["folders", "runs"])
     time.sleep(1)
-    
+
     # Get online ID and retrieve run
     _online_id = _run._local_staging_file.parents[1].joinpath("server_ids", f"{_run._local_staging_file.name.split('.')[0]}.txt").read_text()
     _online_run = Run(_online_id)
-    
+
     assert _online_run.name == _run_name
     assert _online_run.folder == _folder_name
-    
+
     _run.delete()
     _run._local_staging_file.parents[1].joinpath("server_ids", f"{_run._local_staging_file.name.split('.')[0]}.txt").unlink()
     client = Client()
@@ -117,40 +117,51 @@ def test_run_modification_offline() -> None:
     assert _new_run.ttl == 120
     assert _new_run.description == "Simvue test run"
     assert _new_run.name == "simvue_test_run"
-    
+
     sender(_run._local_staging_file.parents[1], 1, 10, ["folders", "runs"])
     time.sleep(1)
-    
+
     # Get online ID and retrieve run
     _online_id = _run._local_staging_file.parents[1].joinpath("server_ids", f"{_run._local_staging_file.name.split('.')[0]}.txt").read_text()
     _online_run = Run(_online_id)
-    
+
     assert _online_run.ttl == 120
     assert _online_run.description == "Simvue test run"
     assert _online_run.name == "simvue_test_run"
     assert _online_run.folder == _folder_name
-    
+
     # Now add a new set of tags in offline mode and send
     _new_run.tags = ["simvue", "test", "tag"]
     _new_run.commit()
-        
+
     # Shouldn't yet be available in the online run since it hasnt been sent
     _online_run.refresh()
     assert _online_run.tags == []
-    
+
     sender(_run._local_staging_file.parents[1], 1, 10, ["folders", "runs"])
     time.sleep(1)
-        
+
     _online_run.refresh()
     assert sorted(_new_run.tags) == sorted(["simvue", "test", "tag"])
     assert sorted(_online_run.tags) == sorted(["simvue", "test", "tag"])
-    
+
     _run.delete()
     _run._local_staging_file.parents[1].joinpath("server_ids", f"{_run._local_staging_file.name.split('.')[0]}.txt").unlink()
     client = Client()
     client.delete_folder(_folder_name, recursive=True, remove_runs=True)
 
 
+@pytest.mark.api
+@pytest.mark.online
+def test_get_run() -> None:
+    _uuid: str = f"{uuid.uuid4()}".split("-")[0]
+    _folder_name = f"/simvue_unit_testing/{_uuid}"
+    _folder = Folder.new(path=_folder_name)
+    _run_1 = Run.new(folder=_folder_name)
+    _run_2 = Run.new(folder=_folder_name)
+    assert len(list(Run._get_all_objects(count=2, offset=None))) == 2
+
+
 @pytest.mark.api
 @pytest.mark.online
 def test_run_get_properties() -> None: