Expect `with_count` when calling `ListDatasetElements`

tests/test_dataset_worker.py

@@ -206,11 +206,26 @@ def test_list_dataset_elements_per_split_api_error(
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
         # The API call is retried 5 times
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
     ]
 
 
@@ -302,7 +317,10 @@ def test_list_dataset_elements_per_split(
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
     ]
 
 

tests/test_elements_worker/test_dataset.py

@@ -144,11 +144,26 @@ def test_list_dataset_elements_api_error(
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
         # The API call is retried 5 times
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
     ]
 
 
@@ -242,7 +257,10 @@ def test_list_dataset_elements(
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/"),
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+        ),
     ]