AIP-81 Improve list operations list_execute

airflow-ctl/src/airflowctl/api/operations.py

@@ -160,26 +160,23 @@ def execute_list(
         params: dict | None = None,
     ) -> T | ServerResponseError:
         shared_params = {**(params or {})}
-        try:
-            self.response = self.client.get(path, params=shared_params)
-            first_pass = data_model.model_validate_json(self.response.content)
-            total_entries = first_pass.total_entries  # type: ignore[attr-defined]
-            if total_entries < limit:
-                return first_pass
-            for key, value in first_pass.model_dump().items():
-                if key != "total_entries" and isinstance(value, list):
-                    break
-            entry_list = getattr(first_pass, key)
+        self.response = self.client.get(path, params=shared_params)
+        first_pass = data_model.model_validate_json(self.response.content)
+        total_entries = first_pass.total_entries  # type: ignore[attr-defined]
+        if total_entries < limit:
+            return first_pass
+        for key, value in first_pass.model_dump().items():
+            if key != "total_entries" and isinstance(value, list):
+                break
+        entry_list = getattr(first_pass, key)
+        offset = offset + limit
+        while offset < total_entries:
+            self.response = self.client.get(path, params={**shared_params, "offset": offset})
+            entry = data_model.model_validate_json(self.response.content)
             offset = offset + limit
-            while offset < total_entries:
-                self.response = self.client.get(path, params={**shared_params, "offset": offset})
-                entry = data_model.model_validate_json(self.response.content)
-                offset = offset + limit
-                entry_list.extend(getattr(entry, key))
-            obj = data_model(**{key: entry_list, "total_entries": total_entries})
-            return data_model.model_validate(obj.model_dump())
-        except ServerResponseError as e:
-            raise e
+            entry_list.extend(getattr(entry, key))
+        obj = data_model(**{key: entry_list, "total_entries": total_entries})
+        return data_model.model_validate(obj.model_dump())
 
 
 # Login operations

airflow-ctl/tests/airflow_ctl/api/test_operations.py

@@ -20,7 +20,9 @@
 import datetime
 import json
 import uuid
+from math import ceil
 from typing import TYPE_CHECKING
+from unittest.mock import Mock
 
 import httpx
 import pytest
@@ -91,6 +93,7 @@
     VariableResponse,
     VersionInfo,
 )
+from airflowctl.api.operations import BaseOperations
 from airflowctl.exceptions import AirflowCtlConnectionException
 
 if TYPE_CHECKING:
@@ -127,7 +130,6 @@ def test_server_connection_refused(self):
     @pytest.mark.parametrize(
         "total_entries, limit, expected_response",
         [
-            (0, 0, (HelloCollectionResponse(hellos=[], total_entries=0))),
             (1, 50, (HelloCollectionResponse(hellos=[HelloResponse(name="hello")], total_entries=1))),
             (
                 150,
@@ -136,33 +138,58 @@ def test_server_connection_refused(self):
                     HelloCollectionResponse(
                         hellos=[
                             HelloResponse(name="hello"),
-                            HelloResponse(name="hello"),
-                            HelloResponse(name="hello"),
-                        ],
+                        ]
+                        * 150,
                         total_entries=150,
                     )
                 ),
             ),
             (
                 90,
                 50,
-                (
-                    HelloCollectionResponse(
-                        hellos=[HelloResponse(name="hello"), HelloResponse(name="hello")], total_entries=90
-                    )
-                ),
+                (HelloCollectionResponse(hellos=[HelloResponse(name="hello")] * 90, total_entries=90)),
             ),
         ],
     )
     def test_execute_list(self, total_entries, limit, expected_response):
-        hello_response = []
-        if total_entries != 0:
-            update = (total_entries + limit - 1) // limit
-            hello_response.extend([HelloResponse(name="hello")] * update)
-        hello_collection_response = HelloCollectionResponse(
-            hellos=hello_response, total_entries=total_entries
+        get_response_mock = []
+
+        mock_client = Mock()
+        mock_client.get.side_effect = get_response_mock
+        base_operation = BaseOperations(client=mock_client)
+
+        nb_of_pages = ceil(total_entries / limit)
+        for page in range(nb_of_pages):
+            if page == nb_of_pages - 1 and (remaining_entries := total_entries % limit) > 0:
+                # partial page
+                get_response_mock.append(
+                    Mock(
+                        content=json.dumps(
+                            {
+                                "hellos": [{"name": "hello"}] * remaining_entries,
+                                "total_entries": total_entries,
+                            }
+                        )
+                    )
+                )
+                continue
+            # page is full
+            get_response_mock.append(
+                Mock(
+                    content=json.dumps(
+                        {
+                            "hellos": [{"name": "hello"}] * limit,
+                            "total_entries": total_entries,
+                        }
+                    )
+                )
+            )
+
+        response = base_operation.execute_list(
+            path="some_fake_path", data_model=HelloCollectionResponse, limit=limit
         )
-        assert expected_response == hello_collection_response
+
+        assert expected_response == response
 
 
 class TestAssetsOperations: