Refactor of c2d receive + test.

Azure · Nov 22, 2018 · 2cbe9c6 · 2cbe9c6
1 parent 6ae09a1
commit 2cbe9c6
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 41 deletions.
diff --git a/azext_iot/device_sdk/iot_hub_gateway_device_apis.py b/azext_iot/device_sdk/iot_hub_gateway_device_apis.py
@@ -170,26 +170,11 @@ def receive_device_bound_notification(
             raise exp
 
         # @digimaun - Altered to support msg headers
-        header_dict = {}
         if response.status_code == 200:
-            header_dict = {
-                'ETag': 'str',
-                'iothub-sequencenumber': 'str',
-                'iothub-enqueuedtime': 'str',
-                'iothub-expiry': 'str',
-                'iothub-deliverycount': 'str',
-                'iothub-messageid': 'str',
-                'iothub-correlationid': 'str',
-                'iothub-userid': 'str',
-                'iothub-to': 'str',
-                'iothub-ack': 'str',
-                'Content-Type': 'str',
-                'Content-Encoding': 'str',
-            }
+            return response
 
         if raw:
             client_raw_response = ClientRawResponse(None, response)
-            client_raw_response.add_headers(header_dict)
             return client_raw_response
     receive_device_bound_notification.metadata = {'url': '/devices/{id}/messages/deviceBound'}
 

diff --git a/azext_iot/operations/hub.py b/azext_iot/operations/hub.py
@@ -935,21 +935,24 @@ def _iot_c2d_message_receive(target, device_id, lock_timeout=60):
         request_headers['IotHub-MessageLockTimeout'] = str(lock_timeout)
 
     try:
-        result = device_sdk.receive_device_bound_notification(device_id, raw=True, custom_headers=request_headers)
-        if result and result.response.status_code == 200:
-            return {
-                'ack': result.headers['iothub-ack'],
-                'correlationId': result.headers['iothub-correlationid'],
-                'data': result.response.text,
-                'deliveryCount': result.headers['iothub-deliverycount'],
-                'enqueuedTime': result.headers['iothub-enqueuedtime'],
-                'expiry': result.headers['iothub-expiry'],
-                'etag': result.headers['ETag'].strip('"'),
-                'messageId': result.headers['iothub-messageid'],
-                'sequenceNumber': result.headers['iothub-sequencenumber'],
-                'to': result.headers['iothub-to'],
-                'userId': result.headers['iothub-userid']
+        result = device_sdk.receive_device_bound_notification(device_id, custom_headers=request_headers)
+        if result and result.status_code == 200:
+            payload = {
+                'ack': result.headers.get('iothub-ack'),
+                'correlationId': result.headers.get('iothub-correlationid'),
+                'data': result.text,
+                'deliveryCount': result.headers.get('iothub-deliverycount'),
+                'enqueuedTime': result.headers.get('iothub-enqueuedtime'),
+                'expiry': result.headers.get('iothub-expiry'),
+                'etag': result.headers.get('ETag'),
+                'messageId': result.headers.get('iothub-messageid'),
+                'sequenceNumber': result.headers.get('iothub-sequencenumber'),
+                'to': result.headers.get('iothub-to'),
+                'userId': result.headers.get('iothub-userid')
             }
+            if payload.get('etag'):
+                payload['etag'] = payload['etag'].strip('"')
+            return payload
         return
     except errors.CloudError as e:
         raise CLIError(unpack_msrest_error(e))

diff --git a/azext_iot/tests/test_iot_ext_unit.py b/azext_iot/tests/test_iot_ext_unit.py
@@ -1565,23 +1565,18 @@ def _build_event(hub):
 
 
 class TestDeviceMessaging():
+    data = '{"data": "value"}'
+
     @pytest.fixture
     def serviceclient(self, mocker, fixture_ghcs, fixture_sas, request):
         service_client = mocker.patch(path_service_client)
-        service_client.return_value = build_mock_response(mocker)
+        service_client.return_value = build_mock_response(mocker, 200, TestDeviceMessaging.data, sample_c2d_receive, raw=True)
         return service_client
 
-    @pytest.fixture
-    def fixture_mocker(self, mocker, fixture_ghcs, fixture_sas, request):
-        return mocker
-
-    def test_c2d_receive(self, fixture_mocker):
-        data = '{"data": "value"}'
-        service_client = fixture_mocker.patch(path_service_client)
-        service_client.return_value = build_mock_response(fixture_mocker, 200, data, sample_c2d_receive, raw=True)
+    def test_c2d_receive(self, serviceclient):
         timeout = 120
         result = subject.iot_c2d_message_receive(fixture_cmd, device_id, mock_target['entity'], timeout)
-        args = service_client.call_args
+        args = serviceclient.call_args
         url = args[0][0].url
         method = args[0][0].method
         headers = args[0][1]
@@ -1592,7 +1587,7 @@ def test_c2d_receive(self, fixture_mocker):
 
         assert result['ack'] == sample_c2d_receive['iothub-ack']
         assert result['correlationId'] == sample_c2d_receive['iothub-correlationid']
-        assert result['data'] == data
+        assert result['data'] == TestDeviceMessaging.data
         assert result['deliveryCount'] == sample_c2d_receive['iothub-deliverycount']
         assert result['enqueuedTime'] == sample_c2d_receive['iothub-enqueuedtime']
         assert result['etag'] == sample_c2d_receive['ETag'].strip('"')