fix: fix list failing without order_by and local sorting (#320)

googleapis · Apr 14, 2021 · 06e99db · 06e99db
1 parent 272fc68
commit 06e99db
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
@@ -923,14 +923,15 @@ def _list_with_local_order(
             credentials=credentials,
         )
 
-        desc = "desc" in order_by
-        order_by = order_by.replace("desc", "")
-        order_by = order_by.split(",")
+        if order_by:
+            desc = "desc" in order_by
+            order_by = order_by.replace("desc", "")
+            order_by = order_by.split(",")
 
-        li.sort(
-            key=lambda x: tuple(getattr(x, field.strip()) for field in order_by),
-            reverse=desc,
-        )
+            li.sort(
+                key=lambda x: tuple(getattr(x, field.strip()) for field in order_by),
+                reverse=desc,
+            )
 
         return li
 

diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py
@@ -849,6 +849,20 @@ def test_list_dataset(self, list_datasets_mock):
         for ds in ds_list:
             assert type(ds) == aiplatform.TabularDataset
 
+    def test_list_dataset_no_order_or_filter(self, list_datasets_mock):
+
+        ds_list = aiplatform.TabularDataset.list()
+
+        list_datasets_mock.assert_called_once_with(
+            request={"parent": _TEST_PARENT, "filter": None}
+        )
+
+        # Ensure returned list is smaller since it filtered out non-tabular datasets
+        assert len(ds_list) < len(_TEST_DATASET_LIST)
+
+        for ds in ds_list:
+            assert type(ds) == aiplatform.TabularDataset
+
 
 class TestTextDataset:
     def setup_method(self):