fix: is_rockpi and is_raspberry_pi should not break outside Balena #69

* read model information from /proc/device-tree/model
in the absense of balena env variable

hm_pyhelper/hardware_definitions.py

@@ -1,30 +1,79 @@
 import os
+from enum import Enum, auto
+from collections import namedtuple
 from hm_pyhelper.exceptions import UnknownVariantException, \
-                                   UnknownVariantAttributeException
+    UnknownVariantAttributeException
 
+BALENA_ENV_RASPBERRY_PI_MODELS = [
+    'raspberry-pi2',
+    'raspberrypi3',
+    'raspberrypi3-64',
+    'raspberrypi4-64',
+    'nebra-hnt',
+    'raspberrypicm4-ioboard'
+]
 
-def is_rockpi():
-    return 'rockpi-4b-rk3399' in os.getenv('BALENA_DEVICE_TYPE')
+BALENA_ENV_ROCKPI_MODELS = ['rockpi-4b-rk3399']
 
+SBCInfo = namedtuple('SBCInfo', ['vendor_id', 'vendor_name', 'model_name'])
 
-def is_raspberry_pi():
+
+class DeviceVendor(Enum):
+    """
+    Enum for device vendors.
+    """
+    INVALID = auto()
+    ROCK_PI = auto()
+    RASPBERRY_PI = auto()
+
+
+def device_model():
+    with open('/proc/device-tree/model', 'r') as f:
+        return f.readline().strip()
+
+
+def sbc_info() -> SBCInfo:
+    '''
+    return SBCInfo formed by reading '/proc/device-tree/model'
+    '''
+    sbc_info = SBCInfo(vendor_id=DeviceVendor.INVALID, vendor_name='', model_name='')
+    dev_model = device_model()
+    if dev_model.lower().find('raspberry') >= 0:
+        sbc_info = SBCInfo(vendor_id=DeviceVendor.RASPBERRY_PI,
+                           vendor_name='Raspberry Pi',
+                           model_name=dev_model)
+    elif dev_model.lower().find('rock') >= 0:
+        sbc_info = SBCInfo(vendor_id=DeviceVendor.ROCK_PI,
+                           vendor_name='Radxa Rock Pi',
+                           model_name=dev_model)
+    return sbc_info
+
+
+def is_rockpi() -> bool:
+    device_type = os.getenv('BALENA_DEVICE_TYPE')
+
+    # use device tree supplied model name if evn not set
+    if not device_type:
+        return sbc_info().vendor_id == DeviceVendor.ROCK_PI
+
+    # honor env override
+    if device_type in BALENA_ENV_ROCKPI_MODELS:
+        return True
+    return False
+
+
+def is_raspberry_pi() -> bool:
     """
     Pulled from
     https://www.balena.io/docs/reference/base-images/devicetypes/
     """
     device_type = os.getenv('BALENA_DEVICE_TYPE')
-    device_type_match = [
-        'raspberry-pi2',
-        'raspberrypi3',
-        'raspberrypi3-64',
-        'raspberrypi4-64',
-        'nebra-hnt',
-        'raspberrypicm4-ioboard'
-    ]
-
-    for device in device_type_match:
-        if device in device_type:
-            return True
+    if not device_type:
+        return sbc_info().vendor_id == DeviceVendor.RASPBERRY_PI
+
+    # honor env override
+    if device_type in BALENA_ENV_RASPBERRY_PI_MODELS:
+        return True
     return False
 
 

hm_pyhelper/tests/test_hardware_definitions.py

@@ -1,12 +1,16 @@
+import os
 import pytest
 from unittest import TestCase
-from unittest.mock import patch
+from unittest.mock import patch, mock_open
 
 from hm_pyhelper.exceptions import UnknownVariantException, \
-                                   UnknownVariantAttributeException
+    UnknownVariantAttributeException
 
-from hm_pyhelper.hardware_definitions import variant_definitions, \
-                                            get_variant_attribute
+from hm_pyhelper.hardware_definitions import is_rockpi, variant_definitions, \
+    get_variant_attribute, is_raspberry_pi, BALENA_ENV_RASPBERRY_PI_MODELS, \
+    BALENA_ENV_ROCKPI_MODELS
+
+BUILTINS_OPEN_LITERAL = "builtins.open"
 
 
 class TestHardwareDefinitions(TestCase):
@@ -73,3 +77,65 @@ def test_get_variant_attribute_unknown_variant(self):
     def test_get_variant_attribute_unknown_attribute(self):
         with pytest.raises(UnknownVariantAttributeException):
             get_variant_attribute('NEBHNT-XYZ', 'Nonexistant')
+
+    # raspberry pi model names picked from pi kernel sources
+    # https://github.com/raspberrypi/linux
+    # grep -ir "raspberry" linux/arch/arm*  | grep "model ="  | cut -d "=" -f2
+    mock_known_dts_pi_models = [
+        "Raspberry Pi Model B+",
+        "Raspberry Pi Model B",
+        "Raspberry Pi Compute Module",
+        "Raspberry Pi Zero",
+        "Raspberry Pi 2 Model B rev 1.2",
+        "Raspberry Pi Compute Module 3",
+        "Raspberry Pi Zero 2 W",
+        "Raspberry Pi 4 Model B",
+        "Raspberry Pi 400",
+        "Raspberry Pi Compute Module 4",
+        "Raspberry Pi Compute Module 4S",
+        "Raspberry Pi Model A+",
+        "Raspberry Pi Model A",
+        "Raspberry Pi Model B rev2",
+        "Raspberry Pi Compute Module IO board rev1",
+        "Raspberry Pi Zero W",
+        "Raspberry Pi 2 Model B",
+        "Raspberry Pi 3 Model A+",
+        "Raspberry Pi 3 Model B+",
+        "Raspberry Pi 3 Model B",
+        "Raspberry Pi Compute Module 3 IO board V3.0"
+    ]
+
+    def test_is_raspberry_pi(self):
+        for model in self.mock_known_dts_pi_models:
+            with patch(BUILTINS_OPEN_LITERAL, new_callable=mock_open, read_data=model):
+                self.assertTrue(is_raspberry_pi())
+            with patch(BUILTINS_OPEN_LITERAL, new_callable=mock_open, read_data="Rock something"):
+                self.assertFalse(is_raspberry_pi())
+
+        # test balena env based detection
+        for model in BALENA_ENV_RASPBERRY_PI_MODELS:
+            with patch.dict(os.environ, {'BALENA_DEVICE_TYPE': model}):
+                self.assertTrue(is_raspberry_pi())
+            # in absence of the env, it should look for /proc/device-tree/model
+            # which will not exist on test environment.
+            with self.assertRaises(FileNotFoundError):
+                self.assertFalse(is_raspberry_pi())
+
+    mock_known_rock_dts_models = ["ROCK PI 4B"]
+
+    def test_is_rock_pi(self):
+        for model in self.mock_known_rock_dts_models:
+            with patch(BUILTINS_OPEN_LITERAL, new_callable=mock_open, read_data=model):
+                self.assertTrue(is_rockpi())
+            with patch(BUILTINS_OPEN_LITERAL, new_callable=mock_open,
+                       read_data="raspberry something"):
+                self.assertFalse(is_rockpi())
+
+        # test balena env based detection
+        for model in BALENA_ENV_ROCKPI_MODELS:
+            with patch.dict(os.environ, {'BALENA_DEVICE_TYPE': model}):
+                self.assertTrue(is_rockpi())
+            # in absence of the env, it should look for /proc/device-tree/model
+            # which will not exist on test environment.
+            with self.assertRaises(FileNotFoundError):
+                self.assertFalse(is_rockpi())

setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name='hm_pyhelper',
-    version='0.13.14',
+    version='0.13.15',
     author="Nebra Ltd",
     author_email="[email protected]",
     description="Helium Python Helper",