Merge pull request #139 from NebraLtd/automatic_sbc_detection

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

hm_pyhelper/hardware_definitions.py

@@ -1,31 +1,14 @@
-import os
+import hm_pyhelper.sbc as sbc
 from hm_pyhelper.exceptions import UnknownVariantException, \
-                                   UnknownVariantAttributeException
+    UnknownVariantAttributeException
 
 
-def is_rockpi():
-    return 'rockpi-4b-rk3399' in os.getenv('BALENA_DEVICE_TYPE')
+def is_rockpi() -> bool:
+    return sbc.is_sbc_type(sbc.DeviceVendorID.ROCK_PI)
 
 
-def is_raspberry_pi():
-    """
-    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
-    return False
+def is_raspberry_pi() -> bool:
+    return sbc.is_sbc_type(sbc.DeviceVendorID.RASPBERRY_PI)
 
 
 variant_definitions = {

hm_pyhelper/sbc.py

@@ -0,0 +1,78 @@
+"""
+This module provides a set of functions to extract information about
+the Single Board Computer in use.
+It considers Balena environment variables as primary source of truth.
+It also uses the device tree to extract information about the SBC.
+"""
+
+import os
+from enum import Enum, auto
+from collections import namedtuple
+
+SBCInfo = namedtuple('SBCInfo', ['vendor_id', 'vendor_name', 'model_name'])
+
+
+class DeviceVendorID(Enum):
+    """
+    Enum for device vendors.
+    """
+    INVALID = auto()
+    ROCK_PI = auto()
+    RASPBERRY_PI = auto()
+
+
+BALENA_ENV_RASPBERRY_PI_MODELS = [
+    """
+    Pulled from
+    https://www.balena.io/docs/reference/base-images/devicetypes/
+    """
+    'raspberry-pi2',
+    'raspberrypi3',
+    'raspberrypi3-64',
+    'raspberrypi4-64',
+    'nebra-hnt',
+    'raspberrypicm4-ioboard'
+]
+
+BALENA_ENV_ROCKPI_MODELS = ['rockpi-4b-rk3399']
+
+BALENA_MODELS = {
+    DeviceVendorID.ROCK_PI: BALENA_ENV_ROCKPI_MODELS,
+    DeviceVendorID.RASPBERRY_PI: BALENA_ENV_RASPBERRY_PI_MODELS
+}
+
+
+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=DeviceVendorID.INVALID, vendor_name='', model_name='')
+    dev_model = device_model()
+    if dev_model.lower().find('raspberry') >= 0:
+        sbc_info = SBCInfo(vendor_id=DeviceVendorID.RASPBERRY_PI,
+                           vendor_name='Raspberry Pi',
+                           model_name=dev_model)
+    elif dev_model.lower().find('rock') >= 0:
+        sbc_info = SBCInfo(vendor_id=DeviceVendorID.ROCK_PI,
+                           vendor_name='Radxa Rock Pi',
+                           model_name=dev_model)
+    return sbc_info
+
+
+def is_sbc_type(device_id: DeviceVendorID) -> bool:
+    '''
+    Return true if the sbc matches the type supplied.
+    '''
+    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 == device_id
+
+    # honor env override
+    return device_type in BALENA_MODELS.get(device_id, [])

hm_pyhelper/tests/test_hardware_definitions.py

@@ -1,12 +1,17 @@
+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
+from hm_pyhelper.sbc import BALENA_ENV_RASPBERRY_PI_MODELS, \
+    BALENA_ENV_ROCKPI_MODELS
+
+BUILTINS_OPEN_LITERAL = "builtins.open"
 
 
 class TestHardwareDefinitions(TestCase):
@@ -73,3 +78,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.15',
+    version='0.13.16',
     author="Nebra Ltd",
     author_email="[email protected]",
     description="Helium Python Helper",