Decouple rte and snipeit modules

Signed-off-by: Pawel Langowski <[email protected]>

Author: PLangowski

osfv_cli/src/osfv/cli/cli.py

@@ -5,6 +5,7 @@
 from copy import copy
 from importlib import metadata
 
+import osfv.libs.utils as utils
 import pexpect
 import requests
 from osfv.libs.rte import RTE
@@ -717,7 +718,8 @@ def main():
                 exit(
                     f"Failed to retrieve model name from Snipe-IT. Check again arguments, or try providing model manually."
                 )
-        rte = RTE(args.rte_ip, dut_model_name, snipeit_api)
+        sonoff, sonoff_ip = utils.init_sonoff(None, args.rte_ip, snipeit_api)
+        rte = RTE(args.rte_ip, dut_model_name, sonoff)
 
         if args.rte_cmd == "rel":
             # Handle RTE relay related commands

osfv_cli/src/osfv/libs/rte.py

@@ -5,7 +5,6 @@
 import yaml
 from importlib_resources import files
 from osfv.libs.rtectrl_api import rtectrl
-from osfv.libs.sonoff_api import SonoffDevice
 from voluptuous import Any, Optional, Required, Schema
 
 
@@ -29,12 +28,15 @@ class RTE(rtectrl):
     PROGRAMMER_CH341A = "ch341a_spi"
     FLASHROM_CMD = "flashrom -p {programmer} {args}"
 
-    def __init__(self, rte_ip, dut_model, snipeit_api=None, sonoff_ip=None):
+    def __init__(self, rte_ip, dut_model, sonoff):
         self.rte_ip = rte_ip
         self.dut_model = dut_model
         self.dut_data = self.load_model_data()
-        self.snipeit_api = snipeit_api
-        self.sonoff, self.sonoff_ip = self.init_sonoff(sonoff_ip)
+        self.sonoff = sonoff
+        if not self.sonoff_sanity_check():
+            raise SonoffNotFound(
+                exit(f"Missing value for 'sonoff_ip' or Sonoff not found in SnipeIT")
+            )
 
     def load_model_data(self):
         file_path = os.path.join(files("osfv"), "models", f"{self.dut_model}.yml")
@@ -97,26 +99,6 @@ def load_model_data(self):
         # Return the loaded data
         return data
 
-    def init_sonoff(self, init_sonoff_ip):
-        sonoff_ip = ""
-        sonoff = None
-
-        if self.dut_data["pwr_ctrl"]["sonoff"] is True:
-            if not self.snipeit_api:
-                if not init_sonoff_ip:
-                    raise TypeError(
-                        f"Expected a value for 'sonoff_ip', but got None"
-                    )
-                sonoff_ip = init_sonoff_ip
-            else:
-                sonoff_ip = self.snipeit_api.get_sonoff_ip_by_rte_ip(self.rte_ip)
-                if not sonoff_ip:
-                    raise SonoffNotFound(
-                        exit(f"Sonoff IP not found in SnipeIT for RTE: {self.rte_ip}")
-                    )
-            sonoff = SonoffDevice(sonoff_ip)
-        return sonoff, sonoff_ip
-
     def power_on(self, sleep=1):
         self.gpio_set(self.GPIO_POWER, "low", sleep)
         time.sleep(sleep)
@@ -322,6 +304,12 @@ def flash_write(self, write_file):
             if self.dut_data["reset_cmos"] == True:
                 self.reset_cmos()
 
+    def sonoff_sanity_check(self):
+        """
+        Verify that if DUT is powered by Sonoff, Sonoff IP is not None
+        """
+        return not self.dut_data["pwr_ctrl"]["sonoff"] or self.sonoff.sonoff_ip
+
 
 class IncompleteModelData(Exception):
     pass

osfv_cli/src/osfv/libs/utils.py

@@ -0,0 +1,12 @@
+from osfv.libs.sonoff_api import SonoffDevice
+
+
+def init_sonoff(init_sonoff_ip, rte_ip, snipeit_api=None):
+    sonoff_ip = ""
+    sonoff = None
+    if not snipeit_api:
+        sonoff_ip = init_sonoff_ip
+    else:
+        sonoff_ip = snipeit_api.get_sonoff_ip_by_rte_ip(rte_ip)
+    sonoff = SonoffDevice(sonoff_ip)
+    return sonoff, sonoff_ip

osfv_cli/src/osfv/rf/rte_robot.py

@@ -1,6 +1,8 @@
+import osfv.libs.utils as utils
 import robot.api.logger
 from osfv.libs.rte import RTE, UnsupportedDUTModel
 from osfv.libs.snipeit_api import SnipeIT
+from osfv.libs.sonoff_api import SonoffDevice
 from robot.api.deco import keyword
 
 model_dict = {
@@ -27,10 +29,11 @@
 
 class RobotRTE:
     def __init__(self, rte_ip, snipeit: bool, sonoff_ip=None, config=None):
+        self.rte_ip = rte_ip
         if snipeit:
-            snipeit_api = SnipeIT()
-            asset_id = snipeit_api.get_asset_id_by_rte_ip(rte_ip)
-            status, dut_model_name = snipeit_api.get_asset_model_name(asset_id)
+            self.snipeit_api = SnipeIT()
+            asset_id = self.snipeit_api.get_asset_id_by_rte_ip(rte_ip)
+            status, dut_model_name = self.snipeit_api.get_asset_model_name(asset_id)
             if status:
                 robot.api.logger.info(
                     f"DUT model retrieved from snipeit: {dut_model_name}"
@@ -39,11 +42,13 @@ def __init__(self, rte_ip, snipeit: bool, sonoff_ip=None, config=None):
                 raise AssertionError(
                     f"Failed to retrieve model name from Snipe-IT. Check again arguments, or try providing model manually."
                 )
-            self.rte = RTE(rte_ip, dut_model_name, snipeit_api)
-        else:
-            self.rte = RTE(
-                rte_ip, self.cli_model_from_osfv(config), sonoff_ip=sonoff_ip
+            self.sonoff, self.sonoff_ip = utils.init_sonoff(
+                sonoff_ip, self.rte_ip, self.snipeit_api
             )
+            self.rte = RTE(rte_ip, dut_model_name, self.sonoff)
+        else:
+            self.sonoff, self.sonoff_ip = utils.init_sonoff(sonoff_ip, self.rte_ip)
+            self.rte = RTE(rte_ip, self.cli_model_from_osfv(config), self.sonoff)
 
     def cli_model_from_osfv(self, osfv_model):
         """