From 433da78f2425af6f32957d8881575f595cd1d62d Mon Sep 17 00:00:00 2001
From: Alexander Allen <arallen@nvidia.com>
Date: Tue, 1 Jun 2021 12:07:33 -0400
Subject: [PATCH] [sonic-py-common] Add platform and chassis info methods to
 device_info (#7652)

#### Why I did it
These methods were added to make some convenient platform and chassis information methods accessible through sonic-py-common. These methods were refactored from sonic-utilities and are used in the `show platform summary` and `show version` commands.

#### How I did it
There are two methods, one is `get_platform_info()` which simply calls local methods to collect useful platform information into a dictionary format, this came directly from sonic-utilities.
---
 .../sonic_py_common/device_info.py            | 48 +++++++++++++++++++
 src/sonic-py-common/tests/device_info_test.py | 10 +++-
 src/sonic-py-common/tests/mock_swsssdk.py     | 16 +++++++
 3 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 src/sonic-py-common/tests/mock_swsssdk.py

diff --git a/src/sonic-py-common/sonic_py_common/device_info.py b/src/sonic-py-common/sonic_py_common/device_info.py
index f486113196fb..2ab20518948e 100644
--- a/src/sonic-py-common/sonic_py_common/device_info.py
+++ b/src/sonic-py-common/sonic_py_common/device_info.py
@@ -32,6 +32,13 @@
 FRONTEND_ASIC_SUB_ROLE = "FrontEnd"
 BACKEND_ASIC_SUB_ROLE = "BackEnd"
 
+# Chassis STATE_DB keys
+CHASSIS_INFO_TABLE = 'CHASSIS_INFO|chassis {}'
+CHASSIS_INFO_CARD_NUM_FIELD = 'module_num'
+CHASSIS_INFO_SERIAL_FIELD = 'serial'
+CHASSIS_INFO_MODEL_FIELD = 'model'
+CHASSIS_INFO_REV_FIELD = 'revision'
+
 
 def get_localhost_info(field):
     try:
@@ -303,6 +310,47 @@ def get_sonic_version_file():
 
     return SONIC_VERSION_YAML_PATH
 
+
+# Get hardware information
+def get_platform_info():
+    """
+    This function is used to get the HW info helper function
+    """
+    from .multi_asic import get_num_asics
+
+    hw_info_dict = {}
+
+    version_info = get_sonic_version_info()
+
+    hw_info_dict['platform'] = get_platform()
+    hw_info_dict['hwsku'] = get_hwsku()
+    hw_info_dict['asic_type'] = version_info['asic_type']
+    hw_info_dict['asic_count'] = get_num_asics()
+
+    return hw_info_dict
+
+
+def get_chassis_info():
+    """
+    This function is used to get the Chassis serial / model / rev number
+    """
+
+    chassis_info_dict = {}
+
+    try:
+        # Init statedb connection
+        db = SonicV2Connector()
+        db.connect(db.STATE_DB)
+        table = CHASSIS_INFO_TABLE.format(1)
+
+        chassis_info_dict['serial'] = db.get(db.STATE_DB, table, CHASSIS_INFO_SERIAL_FIELD)
+        chassis_info_dict['model'] = db.get(db.STATE_DB, table, CHASSIS_INFO_MODEL_FIELD)
+        chassis_info_dict['revision'] = db.get(db.STATE_DB, table, CHASSIS_INFO_REV_FIELD)
+    except Exception:
+        pass
+
+    return chassis_info_dict
+
 #
 # Multi-NPU functionality
 #
diff --git a/src/sonic-py-common/tests/device_info_test.py b/src/sonic-py-common/tests/device_info_test.py
index f3b14b0a5f13..73ac278d46e6 100644
--- a/src/sonic-py-common/tests/device_info_test.py
+++ b/src/sonic-py-common/tests/device_info_test.py
@@ -11,6 +11,7 @@
 
 from sonic_py_common import device_info
 
+from .mock_swsssdk import SonicV2Connector
 
 # TODO: Remove this if/else block once we no longer support Python 2
 if sys.version_info.major == 3:
@@ -49,7 +50,6 @@
     'onie_kernel_version': '4.10.11'
 }
 
-
 class TestDeviceInfo(object):
     @classmethod
     def setup_class(cls):
@@ -70,6 +70,14 @@ def test_get_platform(self):
             result = device_info.get_platform()
             assert result == "x86_64-mlnx_msn2700-r0"
 
+    def test_get_chassis_info(self):
+        with mock.patch("sonic_py_common.device_info.SonicV2Connector", new=SonicV2Connector):
+            result = device_info.get_chassis_info()
+            truth = {"serial": SonicV2Connector.TEST_SERIAL, 
+                     "model": SonicV2Connector.TEST_MODEL, 
+                     "revision": SonicV2Connector.TEST_REV}
+            assert result == truth
+
     @classmethod
     def teardown_class(cls):
         print("TEARDOWN")
diff --git a/src/sonic-py-common/tests/mock_swsssdk.py b/src/sonic-py-common/tests/mock_swsssdk.py
new file mode 100644
index 000000000000..df6affaf1853
--- /dev/null
+++ b/src/sonic-py-common/tests/mock_swsssdk.py
@@ -0,0 +1,16 @@
+class SonicV2Connector:
+    TEST_SERIAL = "MT1822K07815"
+    TEST_MODEL = "MSN2700-CS2FO"
+    TEST_REV = "A1"
+
+    def __init__(self):
+        self.STATE_DB = 'STATE_DB'
+        self.data = {"serial": self.TEST_SERIAL, 
+                     "model": self.TEST_MODEL, 
+                     "revision": self.TEST_REV}
+    
+    def connect(self, db):
+        pass
+
+    def get(self, db, table, field):
+        return self.data.get(field, "N/A")