Handle NotImplementedError exception while changing optoe write max (#2985)

* Handle NotImplementedError exception while changing optoe write max

Signed-off-by: Mihir Patel <[email protected]>

* Added unit test for more coverage

* Removed unused import

---------

Signed-off-by: Mihir Patel <[email protected]>

Author: mihirpat1

sfputil/main.py

@@ -1325,7 +1325,10 @@ def download_firmware(port_name, filepath):
         sys.exit(EXIT_FAIL)
 
     # Increase the optoe driver's write max to speed up firmware download
-    sfp.set_optoe_write_max(SMBUS_BLOCK_WRITE_SIZE)
+    try:
+        sfp.set_optoe_write_max(SMBUS_BLOCK_WRITE_SIZE)
+    except NotImplementedError:
+        click.echo("Platform doesn't implement optoe write max change. Skipping value increase.")
 
     with click.progressbar(length=file_size, label="Downloading ...") as bar:
         address = 0
@@ -1351,7 +1354,10 @@ def download_firmware(port_name, filepath):
             remaining -= count
 
     # Restore the optoe driver's write max to '1' (default value)
-    sfp.set_optoe_write_max(1)
+    try:
+        sfp.set_optoe_write_max(1)
+    except NotImplementedError:
+        click.echo("Platform doesn't implement optoe write max change. Skipping value restore!")
 
     status = api.cdb_firmware_download_complete()
     update_firmware_info_to_state_db(port_name)

tests/sfputil_test.py

@@ -825,6 +825,23 @@ def test_show_fwversion_Rj45(self, mock_chassis):
         assert result.output == 'Show firmware version is not applicable for RJ45 port Ethernet0.\n'
         assert result.exit_code == EXIT_FAIL
 
+    @patch('builtins.open')
+    @patch('sfputil.main.platform_chassis')
+    @patch('sfputil.main.logical_port_to_physical_port_index', MagicMock(return_value=1))
+    @patch('sfputil.main.update_firmware_info_to_state_db', MagicMock())
+    def test_download_firmware(self, mock_chassis, mock_file):
+        mock_file.return_value.tell.return_value = 0
+        mock_sfp = MagicMock()
+        mock_api = MagicMock()
+        mock_sfp.get_xcvr_api = MagicMock(return_value=mock_api)
+        mock_chassis.get_sfp = MagicMock(return_value=mock_sfp)
+        mock_api.get_module_fw_mgmt_feature.return_value = {'status': True, 'feature': (0, 0, False, False, 0)}
+        mock_api.cdb_start_firmware_download.return_value = 1
+        mock_api.cdb_firmware_download_complete.return_value = 1
+        mock_sfp.set_optoe_write_max = MagicMock(side_effect=NotImplementedError)
+        status = sfputil.download_firmware("Ethernet0", "test.bin")
+        assert status == 1
+
     @patch('sfputil.main.platform_chassis')
     @patch('sfputil.main.logical_port_to_physical_port_index', MagicMock(return_value=1))
     def test_run_firmwre(self, mock_chassis):