From 19b03142cc933cf10fdffbe470b60257ac6a56db Mon Sep 17 00:00:00 2001
From: Alexander Allen <arallen@nvidia.com>
Date: Wed, 12 May 2021 23:30:47 +0000
Subject: [PATCH] Add unit test for get_chassis_info

---
 src/sonic-py-common/tests/device_info_test.py | 10 +++++++++-
 src/sonic-py-common/tests/mock_swsssdk.py     | 13 +++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 src/sonic-py-common/tests/mock_swsssdk.py

diff --git a/src/sonic-py-common/tests/device_info_test.py b/src/sonic-py-common/tests/device_info_test.py
index f3b14b0a5f13..eee563c255ec 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 . import mock_swsssdk
 
 # 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("device_info.swsssdk", new=mock_swsssdk):
+            result = device_info.get_chassis_info()
+            truth = {"serial": self.TEST_SERIAL, 
+                     "model": self.TEST_MODEL, 
+                     "revision": self.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..c88e945e09c5
--- /dev/null
+++ b/src/sonic-py-common/tests/mock_swsssdk.py
@@ -0,0 +1,13 @@
+class SonicV2Connector:
+
+    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):
+        self.data.get(field, "N/A")