fix: add paging to az iot dps enrollment-group registration list (A…

…zure#708)

HISTORY.rst

@@ -4,6 +4,14 @@ Release History
 ===============
 
 
+0.24.1
++++++++++++++++
+
+**DPS updates**
+
+* Fix for `az iot dps enrollement-group registration list` to support paging.
+
+
 0.24.0
 +++++++++++++++
 

azext_iot/constants.py

@@ -7,7 +7,7 @@
 
 import os
 
-VERSION = "0.24.0"
+VERSION = "0.24.1"
 EXTENSION_NAME = "azure-iot"
 EXTENSION_ROOT = os.path.dirname(os.path.abspath(__file__))
 EXTENSION_CONFIG_ROOT_KEY = "iotext"

azext_iot/operations/dps.py

@@ -816,10 +816,14 @@ def _get_dps_connection_string(
 
 
 # DPS Registration
-
-
 def iot_dps_registration_list(
-    cmd, enrollment_id, dps_name=None, resource_group_name=None, login=None, auth_type_dataplane=None,
+    cmd,
+    enrollment_id,
+    dps_name=None,
+    resource_group_name=None,
+    top=None,
+    login=None,
+    auth_type_dataplane=None,
 ):
     discovery = DPSDiscovery(cmd)
     target = discovery.get_target(
@@ -831,10 +835,7 @@ def iot_dps_registration_list(
     try:
         resolver = SdkResolver(target=target)
         sdk = resolver.get_sdk(SdkType.dps_sdk)
-
-        return sdk.device_registration_state.query(
-            enrollment_id, raw=True
-        ).response.json()
+        return _execute_query([enrollment_id], sdk.device_registration_state.query, top)
     except ProvisioningServiceErrorDetailsException as e:
         handle_service_exception(e)
 

azext_iot/tests/dps/device_registration/test_iot_device_registration_group_int.py

@@ -195,6 +195,15 @@ def test_dps_device_registration_symmetrickey_lifecycle(provisioned_iot_dps_modu
         assert registration["status"] == "assigned"
 
         # Check for both registration from service side
+        random_registration = cli.invoke(
+            set_cmd_auth_type(
+                f"iot dps enrollment-group registration list --dps-name {dps_name} -g {dps_rg} --group-id {group_id} --top 1",
+                auth_type=auth_phase,
+                cstring=dps_cstring
+            ),
+        ).as_json()
+        assert len(random_registration) == 1
+
         service_side_registrations = cli.invoke(
             set_cmd_auth_type(
                 f"iot dps enrollment-group registration list --dps-name {dps_name} -g {dps_rg} --group-id {group_id}",

azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_unit.py

@@ -721,25 +721,82 @@ def serviceclient(self, mocked_response, fixture_gdcs, fixture_dps_sas, patch_ce
         mocked_response.add(
             method=responses.POST,
             url="https://{}/registrations/{}/query?".format(mock_dps_target['entity'], enrollment_id),
-            body=json.dumps([generate_registration_state_show()]),
+            body=json.dumps([
+                generate_registration_state_show(),
+                generate_registration_state_show(),
+                generate_registration_state_show(),
+                generate_registration_state_show()
+            ]),
             status=request.param,
             content_type="application/json",
             match_querystring=False
         )
         yield mocked_response
 
-    def test_registration_list(self, serviceclient, fixture_cmd):
-        subject.iot_dps_registration_list(
+    @pytest.mark.parametrize("top", [None, 3])
+    def test_registration_list(self, serviceclient, fixture_cmd, top):
+        result = subject.iot_dps_registration_list(
             cmd=fixture_cmd,
             dps_name=mock_dps_target['entity'],
             enrollment_id=enrollment_id,
             resource_group_name=resource_group,
+            top=top
         )
         request = serviceclient.calls[0].request
         url = request.url
         method = request.method
         assert "{}/registrations/{}/query?".format(mock_dps_target['entity'], enrollment_id) in url
         assert method == 'POST'
+        if top:
+            assert len(result) == top
+
+    @pytest.fixture(params=[200])
+    def pagingserviceclient(
+        self, mocked_response, fixture_gdcs, fixture_dps_sas, patch_certificate_open, request
+    ):
+        mocked_response.assert_all_requests_are_fired = False
+        mocked_response.add(
+            method=responses.POST,
+            url="https://{}/registrations/{}/query?".format(mock_dps_target['entity'], enrollment_id),
+            body=json.dumps([
+                generate_registration_state_show(),
+                generate_registration_state_show(),
+                generate_registration_state_show(),
+                generate_registration_state_show()
+            ]),
+            headers={"x-ms-continuation": "continuation_token123"},
+            status=request.param,
+            content_type="application/json",
+            match_querystring=False
+        )
+        mocked_response.add(
+            method=responses.POST,
+            url="https://{}/registrations/{}/query?".format(mock_dps_target['entity'], enrollment_id),
+            body=json.dumps([
+                generate_registration_state_show(),
+                generate_registration_state_show(),
+                generate_registration_state_show(),
+            ]),
+            status=request.param,
+            content_type="application/json",
+            match_querystring=False
+        )
+        yield mocked_response
+
+    @pytest.mark.parametrize("top", [None, 3, 6])
+    def test_registration_list_paging(self, pagingserviceclient, fixture_cmd, top):
+        result = subject.iot_dps_registration_list(
+            cmd=fixture_cmd,
+            dps_name=mock_dps_target['entity'],
+            enrollment_id=enrollment_id,
+            resource_group_name=resource_group,
+            top=top
+        )
+        if top:
+            assert len(result) == top
+            assert len(pagingserviceclient.calls) == (1 + top // 4)
+        else:
+            assert len(pagingserviceclient.calls) == 2
 
     def test_registration_list_error(self, fixture_cmd):
         with pytest.raises(CLIError):