Added fix for checksum verification and blackbox tests for firmware update

Author: imbeacon

examples/device/firmware_update.py

@@ -16,7 +16,6 @@
 
 import asyncio
 import logging
-from time import monotonic
 
 from tb_mqtt_client.common.config_loader import DeviceConfig
 from tb_mqtt_client.common.logging_utils import configure_logging, get_logger
@@ -30,25 +29,24 @@
 firmware_received = asyncio.Event()
 firmware_update_timeout = 30
 
+config = DeviceConfig()
+config.host = "localhost"
+config.access_token = "YOUR_ACCESS_TOKEN"
 
-async def firmware_update_callback(_, payload):
-    logger.info(f"Firmware update payload received: {payload}")
+
+async def firmware_update_callback(firmware_data, firmware_info):
+    logger.info(f"Firmware update payload received: {firmware_info}")
     firmware_received.set()
 
 
 async def main():
-    config = DeviceConfig()
-    config.host = "localhost"
-    config.access_token = "YOUR_ACCESS_TOKEN"
 
     client = DeviceClient(config)
     await client.connect()
 
-    await client.update_firmware(on_received_callback=firmware_update_callback)
+    await client.update_firmware(on_received_callback=firmware_update_callback, save_firmware=False)  # Set save_firmware to True if you want to save the firmware data
 
-    update_started = monotonic()
-    while not firmware_received.is_set() and monotonic() - update_started < firmware_update_timeout:
-        await asyncio.sleep(1)
+    await asyncio.wait_for(firmware_received.wait(), timeout=firmware_update_timeout)
 
     await client.stop()
 

tb_mqtt_client/service/device/client.py

@@ -106,7 +106,7 @@ def __init__(self, config: Optional[Union[DeviceConfig, Dict]] = None):
 
         self._firmware_updater = FirmwareUpdater(self)
 
-    async def update_firmware(self, on_received_callback: Optional[Callable[[str], Awaitable[None]]] = None,
+    async def update_firmware(self, on_received_callback: Optional[Callable[[bytes, dict], Awaitable[None]]] = None,
                               save_firmware: bool = True, firmware_save_path: Optional[str] = None):
         await self._firmware_updater.update(on_received_callback, save_firmware, firmware_save_path)
 

tb_mqtt_client/service/device/firmware_updater.py

@@ -102,7 +102,7 @@ async def _get_next_chunk(self):
 
         topic = mqtt_topics.build_firmware_update_request_topic(self._firmware_request_id, self._current_chunk)
         mqtt_message = MqttPublishMessage(topic, payload)
-        await self._client._message_queue.publish(mqtt_message, wait_for_publish=True)
+        await self._client._message_queue.publish(mqtt_message)
 
     async def _verify_downloaded_firmware(self):
         self._log.info('Verifying downloaded firmware...')
@@ -111,8 +111,8 @@ async def _verify_downloaded_firmware(self):
         await self._send_current_firmware_info()
 
         verified = self.verify_checksum(self._firmware_data,
-                                        self._target_checksum,
-                                        self._target_checksum_alg)
+                                        self._target_checksum_alg,
+                                        self._target_checksum)
 
         if verified:
             self._log.debug('Checksum verified.')
@@ -201,8 +201,8 @@ async def _firmware_info_callback(self, response, *args, **kwargs):
 
                 self._firmware_request_id += 1
                 self._target_firmware_length = fetched_firmware_info[FW_SIZE_ATTR]
-                self._target_checksum = fetched_firmware_info[FW_CHECKSUM_ALG_ATTR]
-                self._target_checksum_alg = fetched_firmware_info[FW_CHECKSUM_ATTR]
+                self._target_checksum = fetched_firmware_info[FW_CHECKSUM_ATTR]
+                self._target_checksum_alg = fetched_firmware_info[FW_CHECKSUM_ALG_ATTR]
                 self._target_title = fetched_firmware_info[FW_TITLE_ATTR]
                 self._target_version = fetched_firmware_info[FW_VERSION_ATTR]
 
@@ -239,27 +239,28 @@ def verify_checksum(self, firmware_data, checksum_alg, checksum):
         checksum_of_received_firmware = None
 
         self._log.debug('Checksum algorithm is: %s' % checksum_alg)
-        if checksum_alg.lower() == "sha256":
+        lower_checksum_alg = checksum_alg.lower()
+        if lower_checksum_alg == "sha256":
             checksum_of_received_firmware = sha256(firmware_data).digest().hex()
-        elif checksum_alg.lower() == "sha384":
+        elif lower_checksum_alg == "sha384":
             checksum_of_received_firmware = sha384(firmware_data).digest().hex()
-        elif checksum_alg.lower() == "sha512":
+        elif lower_checksum_alg == "sha512":
             checksum_of_received_firmware = sha512(firmware_data).digest().hex()
-        elif checksum_alg.lower() == "md5":
+        elif lower_checksum_alg == "md5":
             checksum_of_received_firmware = md5(firmware_data).digest().hex()
-        elif checksum_alg.lower() == "murmur3_32":
+        elif lower_checksum_alg == "murmur3_32":
             reversed_checksum = f'{hash(firmware_data, signed=False):0>2X}'
             if len(reversed_checksum) % 2 != 0:
                 reversed_checksum = '0' + reversed_checksum
             checksum_of_received_firmware = "".join(
                 reversed([reversed_checksum[i:i + 2] for i in range(0, len(reversed_checksum), 2)])).lower()
-        elif checksum_alg.lower() == "murmur3_128":
+        elif lower_checksum_alg == "murmur3_128":
             reversed_checksum = f'{hash128(firmware_data, signed=False):0>2X}'
             if len(reversed_checksum) % 2 != 0:
                 reversed_checksum = '0' + reversed_checksum
             checksum_of_received_firmware = "".join(
                 reversed([reversed_checksum[i:i + 2] for i in range(0, len(reversed_checksum), 2)])).lower()
-        elif checksum_alg.lower() == "crc32":
+        elif lower_checksum_alg == "crc32":
             reversed_checksum = f'{crc32(firmware_data) & 0xffffffff:0>2X}'
             if len(reversed_checksum) % 2 != 0:
                 reversed_checksum = '0' + reversed_checksum
@@ -268,11 +269,4 @@ def verify_checksum(self, firmware_data, checksum_alg, checksum):
         else:
             self._log.error('Client error. Unsupported checksum algorithm.')
 
-        self._log.debug(checksum_of_received_firmware)
-
-        random_value = randint(0, 5)
-        if random_value > 3:
-            self._log.debug('Dummy fail! Do not panic, just restart and try again the chance of this fail is ~20%')
-            return False
-
         return checksum_of_received_firmware == checksum

tests/blackbox/conftest.py

@@ -18,12 +18,14 @@
 import pytest
 import requests
 import logging
+import uuid
 
 from requests import HTTPError
 
 from tb_mqtt_client.common.config_loader import GatewayConfig
 from tb_mqtt_client.common.config_loader import DeviceConfig
-from tests.blackbox.rest_helpers import find_related_entity_id, get_device_info_by_id
+from tests.blackbox.rest_helpers import find_related_entity_id, get_device_info_by_id, \
+    create_device_profile_and_firmware
 
 TB_HOST = os.getenv("SDK_BLACKBOX_TB_HOST", "localhost")
 TB_HTTP_PORT = int(os.getenv("SDK_BLACKBOX_TB_PORT", 8080))
@@ -226,3 +228,42 @@ def gateway_config(gateway_info, tb_admin_headers):
 
     requests.delete(f"{TB_URL}/api/deviceProfile/{sub_device_info['deviceProfileId']['id']}",
                     headers=tb_admin_headers)
+
+
+@pytest.fixture
+def firmware_profile_and_package(tb_admin_headers, test_config):
+    firmware_name = 'pytest-firmware'
+    firmware_version = '1.0.0'
+    boundary = uuid.uuid4().hex
+    firmware_bytes = b'Firmware binary data for pytest'
+
+    multipart_body = (
+                         f"--{boundary}\r\n"
+                         f"Content-Disposition: form-data; name=\"file\"; filename=\"firmware.bin\"\r\n"
+                         f"Content-Type: application/octet-stream\r\n\r\n"
+                     ).encode() + firmware_bytes + f"\r\n--{boundary}--\r\n".encode()
+
+    firmware_info = {
+        "name": firmware_name,
+        "version": firmware_version,
+        "data": firmware_bytes
+    }
+
+    firmware_headers = {
+        "Content-Type": f"multipart/form-data; boundary={boundary}",
+        "Content-Length": str(len(multipart_body))
+    }
+    firmware_headers["X-Authorization"] = tb_admin_headers["X-Authorization"]
+
+    device_profile, firmware = create_device_profile_and_firmware(
+        firmware_name, firmware_version, multipart_body, test_config['tb_url'], tb_admin_headers, firmware_headers
+    )
+
+    assert device_profile is not None, "Device profile should be created successfully"
+    assert firmware is not None, "Firmware should be created successfully"
+
+    yield device_profile, firmware, firmware_info
+
+    # Cleanup
+    requests.delete(f"{test_config['tb_url']}/api/deviceProfile/{device_profile['id']['id']}", headers=tb_admin_headers)
+    requests.delete(f"{test_config['tb_url']}/api/firmware/{firmware['id']['id']}", headers=tb_admin_headers)

tests/blackbox/rest_helpers.py

@@ -14,6 +14,8 @@
 
 import asyncio
 from time import time
+from typing import Tuple
+
 import requests
 
 
@@ -134,6 +136,78 @@ def create_device_profile_with_provisioning(device_profile_name: str,
     else:
         response.raise_for_status()
 
+def create_device_profile_and_firmware(firmware_name: str,
+                                       firmware_version: str,
+                                       firmware_data: bytes,
+                                       base_url: str,
+                                       headers: dict,
+                                       firmware_data_headers: dict) -> Tuple[dict, dict]:
+    device_profile = get_default_device_profile(base_url, headers)
+    device_profile['id'] = None
+    device_profile['createdTime'] = None
+    device_profile['name'] = 'pytest-firmware-profile'
+    if 'profileData' not in device_profile:
+        device_profile['profileData'] = {
+                                            "configuration": {
+                                                "type": "DEFAULT"
+                                            },
+                                            "transportConfiguration": {
+                                                "type": "DEFAULT"
+                                            }
+        }
+    response = requests.post(f"{base_url}/api/deviceProfile", json=device_profile, headers=headers)
+    if response.status_code == 200:
+        device_profile = response.json()
+    else:
+        response.raise_for_status()
+
+    # Create OTA package
+    init_ota_package ={
+        "id": None,
+        "createdTime": None,
+        "deviceProfileId": {
+            "entityType": "DEVICE_PROFILE",
+            "id": device_profile['id']['id']
+        },
+        "type": "FIRMWARE",
+        "title": firmware_name,
+        "version": firmware_version,
+        "tag": firmware_name + " " + firmware_version,
+        "url": None,
+        "hasData": False,
+        "fileName": None,
+        "contentType": None,
+        "checksumAlgorithm": None,
+        "checksum": None,
+        "dataSize": None,
+        "externalId": None,
+        "name": firmware_name,
+        "additionalInfo": {
+            "description": ""
+        }
+    }
+    created_ota_package = requests.post(f"{base_url}/api/otaPackage", json=init_ota_package, headers=headers)
+    initial_ota_package = None
+    if created_ota_package.status_code == 200:
+        initial_ota_package = created_ota_package.json()
+    else:
+        created_ota_package.raise_for_status()
+    # Upload firmware data
+    firmware_data_url = f"{base_url}/api/otaPackage/{initial_ota_package['id']['id']}?checksumAlgorithm=SHA256"
+    response = requests.post(firmware_data_url, data=firmware_data, headers=firmware_data_headers)
+    if response.status_code == 200:
+        return device_profile, response.json()
+    else:
+        response.raise_for_status()
+
+def save_device(device: dict, base_url: str, headers: dict) -> dict:
+    url = f"{base_url}/api/device"
+    response = requests.post(url, json=device, headers=headers)
+    if response.status_code == 200:
+        return response.json()
+    else:
+        response.raise_for_status()
+
 async def send_rpc_request(device_id: str, base_url: str, headers: dict, request: dict):
     loop = asyncio.get_running_loop()
 
@@ -144,4 +218,4 @@ def _send():
         return r.json()
 
     result = await loop.run_in_executor(None, _send)
-    return result
+    return result

tests/blackbox/test_client_provisioning.py renamed to tests/blackbox/test_device_client_provisioning.py

@@ -0,0 +1,64 @@
+#  Copyright 2025 ThingsBoard
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import asyncio
+from copy import deepcopy
+
+import pytest
+
+from examples.device import firmware_update
+from tests.blackbox.rest_helpers import save_device
+
+
+@pytest.mark.asyncio
+@pytest.mark.blackbox
+async def test_device_firmware_update(firmware_profile_and_package, device_info, device_config, tb_admin_headers, test_config):
+
+    handler_future = asyncio.Future()
+
+    async def firmware_update_callback(firmware_data, current_firmware_info):
+        handler_future.set_result((firmware_data, current_firmware_info))
+
+    firmware_update.config = device_config
+    firmware_update.firmware_update_callback = firmware_update_callback
+
+    device_profile, firmware_package, firmware_info = firmware_profile_and_package
+
+    device = deepcopy(device_info)
+
+    device['deviceProfileId'] = device_profile['id']
+    device['firmwareId'] = firmware_package['id']
+
+    save_device(device, test_config['tb_url'], tb_admin_headers)
+
+    await asyncio.sleep(1)  # Ensure the device is saved before starting the firmware update
+
+    task = asyncio.create_task(firmware_update.main())
+    result = None
+    try:
+        result = await asyncio.wait_for(handler_future, timeout=30)
+    except Exception as e:
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
+    assert result is not None, "Firmware update callback should be called"
+    assert isinstance(result, tuple), "Result should be a tuple containing firmware data and info"
+    received_firmware_data, received_firmware_info = result
+    assert received_firmware_info['current_fw_title'] == firmware_info['name']
+    assert received_firmware_info['current_fw_version'] == firmware_info['version'], \
+        "Received firmware version should match the package version"
+    assert received_firmware_info['fw_state'] == 'UPDATED'

tests/blackbox/test_device_firmware_update.py

@@ -259,9 +259,3 @@ def test_verify_checksum_known_algorithms(updater, alg):
     name, checksum = alg
     with patch("tb_mqtt_client.service.device.firmware_updater.randint", return_value=0):
         assert updater.verify_checksum(b"data", name, checksum) is True
-
-
-def test_verify_checksum_random_failure(updater):
-    with patch("tb_mqtt_client.service.device.firmware_updater.randint", return_value=5):
-        result = updater.verify_checksum(b"data", "md5", md5(b"data").digest().hex())
-        assert not result