rte: add PSU commands and documentation

osfv_cli/src/osfv/cli/cli.py

@@ -249,6 +249,12 @@ def relay_get(rte, args):
 
 
 def power_on(rte, args):
+    state = rte.psu_get()
+    if state != rte.PSU_STATE_ON:
+        print(f"Power supply state: {state} !")
+        print(
+            'If you wanted to power on the DUT, you need to enable power supply first ("pwr psu on"), pushing the power button is not enough!'
+        )
     print(f"Powering on...")
     rte.power_on(args.time)
 
@@ -263,6 +269,21 @@ def reset(rte, args):
     rte.reset(args.time)
 
 
+def psu_on(rte, args):
+    print(f"Enabling power supply...")
+    rte.psu_on()
+
+
+def psu_off(rte, args):
+    print(f"Disabling power supply...")
+    rte.psu_off()
+
+
+def psu_get(rte, args):
+    state = rte.psu_get()
+    print(f"Power supply state: {state}")
+
+
 def gpio_get(rte, args):
     state = rte.gpio_get(args.gpio_no)
     print(f"GPIO {args.gpio_no} state: {state}")
@@ -663,27 +684,42 @@ def main():
 
     # Power subcommands
     pwr_subparsers = pwr_parser.add_subparsers(title="subcommands", dest="pwr_cmd")
-    power_on_parser = pwr_subparsers.add_parser("on", help="Power on")
+    power_on_parser = pwr_subparsers.add_parser(
+        "on", help="Short power button press, to power on DUT"
+    )
     power_on_parser.add_argument(
         "--time",
         type=int,
         default=1,
         help="Power button press time in seconds (default: 1)",
     )
-    power_off_parser = pwr_subparsers.add_parser("off", help="Power off")
+    power_off_parser = pwr_subparsers.add_parser(
+        "off", help="Long power button press, to power off DUT"
+    )
     power_off_parser.add_argument(
         "--time",
         type=int,
-        default=5,
-        help="Power button press time in seconds (default: 5)",
+        default=6,
+        help="Power button press time in seconds (default: 6)",
+    )
+    reset_parser = pwr_subparsers.add_parser(
+        "reset", help="Reset button press, to reset DUT"
     )
-    reset_parser = pwr_subparsers.add_parser("reset", help="Reset")
     reset_parser.add_argument(
         "--time",
         type=int,
         default=1,
         help="Reset button press time in seconds (default: 1)",
     )
+    psu_parser = pwr_subparsers.add_parser(
+        "psu", help="Generic control interface of the power supply"
+    )
+    psu_subparsers = psu_parser.add_subparsers(
+        title="Power supply commands", dest="psu_cmd"
+    )
+    psu_subparsers.add_parser("on", help="Turn the power supply on")
+    psu_subparsers.add_parser("off", help="Turn the power supply off")
+    psu_subparsers.add_parser("get", help="Display information on DUT's power state")
 
     # GPIO subcommands
     gpio_subparsers = gpio_parser.add_subparsers(title="subcommands", dest="gpio_cmd")
@@ -843,6 +879,13 @@ def main():
                 power_off(rte, args)
             elif args.pwr_cmd == "reset":
                 reset(rte, args)
+            elif args.pwr_cmd == "psu":
+                if args.psu_cmd == "on":
+                    psu_on(rte, args)
+                elif args.psu_cmd == "off":
+                    psu_off(rte, args)
+                elif args.psu_cmd == "get":
+                    psu_get(rte, args)
         elif args.rte_cmd == "serial":
             open_dut_serial(rte, args)
         elif args.rte_cmd == "spi":

osfv_cli/src/osfv/libs/rte.py

@@ -20,6 +20,9 @@ class RTE(rtectrl):
 
     GPIO_CMOS = 12
 
+    PSU_STATE_ON = "ON"
+    PSU_STATE_OFF = "OFF"
+
     SSH_USER = "root"
     SSH_PWD = "meta-rte"
     FW_PATH_WRITE = "/data/write.rom"
@@ -101,40 +104,84 @@ def load_model_data(self):
         return data
 
     def power_on(self, sleep=1):
+        """
+        Turns the power on by setting the power button pin to "low" for a
+        specified duration.
+
+        Args:
+            sleep (int, optional): The duration (in seconds) to keep the
+            power pin in the "low" state. Default is 1 second.
+        """
         self.gpio_set(self.GPIO_POWER, "low", sleep)
         time.sleep(sleep)
 
     def power_off(self, sleep=6):
+        """
+        Turns the power off by setting the power button pin to "low" for a
+        specified duration.
+
+        Args:
+            sleep (int, optional): The duration (in seconds) to keep the power
+            pin in the "low" state. Default is 6 seconds.
+        """
         self.gpio_set(self.GPIO_POWER, "low", sleep)
         time.sleep(sleep)
 
     def reset(self, sleep=1):
+        """
+        Resets the system by setting the reset button pin to "low" for a
+        specified duration.
+
+        Args:
+            sleep (int, optional): The duration (in seconds) to keep the reset
+            pin in the "low" state. Default is 1 second.
+        """
         self.gpio_set(self.GPIO_RESET, "low", sleep)
         time.sleep(sleep)
 
     def relay_get(self):
+        """
+        Retrieves the current state of the relay.
+
+        Returns:
+            str: The state of the relay, either "on" if the GPIO relay pin is
+                 "high", or "off" if the GPIO relay pin is "low".
+        """
         gpio_state = self.gpio_get(self.GPIO_RELAY)
         relay_state = None
         if gpio_state == "high":
-            relay_state = "on"
+            relay_state = self.PSU_STATE_ON
         if gpio_state == "low":
-            relay_state = "off"
+            relay_state = self.PSU_STATE_OFF
         return relay_state
 
     def relay_set(self, relay_state):
+        """
+        Sets the state of the relay by configuring the GPIO relay pin accordingly.
+
+        Args:
+            relay_state (str): Desired state of the relay, either "on" (sets GPIO pin to "high")
+                               or "off" (sets GPIO pin to "low").
+        """
         gpio_state = None
-        if relay_state == "on":
+        if relay_state == self.PSU_STATE_ON:
             gpio_state = "high"
-        if relay_state == "off":
+        if relay_state == self.PSU_STATE_OFF:
             gpio_state = "low"
         self.gpio_set(self.GPIO_RELAY, gpio_state)
 
     def reset_cmos(self):
+        """
+        Resets the CMOS by setting the GPIO CMOS pin to "low" for 10 seconds, then returning it to a "high-z" state.
+        """
         self.gpio_set(self.GPIO_CMOS, "low")
         time.sleep(10)
         self.gpio_set(self.GPIO_CMOS, "high-z")
 
     def spi_enable(self):
+        """
+        Enables the SPI interface by configuring GPIO pins based on the voltage level required by the flash chip.
+        """
         voltage = self.dut_data["flash_chip"]["voltage"]
 
         if voltage == "1.8V":
@@ -152,10 +199,16 @@ def spi_enable(self):
         time.sleep(10)
 
     def spi_disable(self):
+        """
+        Disables the SPI interface by setting the GPIO pins to a "high-z" state.
+        """
         self.gpio_set(self.GPIO_SPI_VCC, "high-z")
         self.gpio_set(self.GPIO_SPI_ON, "high-z")
 
-    def pwr_ctrl_on(self):
+    def psu_on(self):
+        """
+        Connect main power supply to the DUT by setting either relay or Sonoff to ON state.
+        """
         if self.dut_data["pwr_ctrl"]["sonoff"] is True:
             self.sonoff.turn_on()
             state = self.sonoff.get_state()
@@ -168,19 +221,34 @@ def pwr_ctrl_on(self):
                 raise Exception("Failed to power control ON")
         time.sleep(5)
 
-    def pwr_ctrl_off(self):
+    def psu_off(self):
+        """
+        Disconnect main power supply from the DUT by setting either relay or Sonoff to OFF state.
+        """
+        # TODO: rework using abstract interfaces for power control?
         if self.dut_data["pwr_ctrl"]["sonoff"] is True:
             self.sonoff.turn_off()
             state = self.sonoff.get_state()
-            if state != "OFF":
+            if state != self.PSU_STATE_OFF:
                 raise Exception("Failed to power control OFF")
         elif self.dut_data["pwr_ctrl"]["relay"] is True:
-            self.relay_set("off")
+            self.relay_set(self.PSU_STATE_OFF)
             state = self.relay_get()
-            if state != "off":
+            if state != self.PSU_STATE_OFF:
                 raise Exception("Failed to power control OFF")
         time.sleep(2)
 
+    def psu_get(self):
+        """
+        Get PSU state
+        """
+        state = None
+        if self.dut_data["pwr_ctrl"]["sonoff"] is True:
+            state = self.sonoff.get_state()
+        elif self.dut_data["pwr_ctrl"]["relay"] is True:
+            state = self.relay_get()
+        return state
+
     def discharge_psu(self):
         """
         Push power button 5 times in the loop to make sure the charge from PSU is dissipated
@@ -189,9 +257,12 @@ def discharge_psu(self):
             self.power_off(3)
 
     def pwr_ctrl_before_flash(self, programmer, power_state):
+        """
+        Move the DUT into specific power state required for external flashing operation. Defined in the model config file.
+        """
 
         # Always start from the same state (PSU active)
-        self.pwr_ctrl_on()
+        self.psu_on()
         time.sleep(5)
         # Put the device into S5 state
         self.power_off(6)
@@ -210,19 +281,25 @@ def pwr_ctrl_before_flash(self, programmer, power_state):
             pass
         elif power_state == "G3":
             # Turn off the PSU/AC brick to put device into G3
-            self.pwr_ctrl_off()
+            self.psu_off()
             self.discharge_psu()
         else:
             exit(
                 f"Power state: '{power_state}' is not supported. Please check model config."
             )
 
     def pwr_ctrl_after_flash(self, programmer):
+        """
+        Additional power actions to take after flashing.
+        """
         if programmer == "rte_1_1":
             self.spi_disable()
             time.sleep(2)
 
     def flash_cmd(self, args, read_file=None, write_file=None):
+        """
+        Send the firmware file to RTE and execute flashrom command over SSH to flash the DUT.
+        """
         try:
             self.pwr_ctrl_before_flash(
                 self.dut_data["programmer"]["name"],
@@ -307,6 +384,9 @@ def flash_cmd(self, args, read_file=None, write_file=None):
         return flashrom_rc
 
     def flash_create_args(self, extra_args=""):
+        """
+        Create flashrom arguments based on the DUT model config.
+        """
         args = ""
 
         # Set chip explicitly, if defined in model configuration
@@ -320,18 +400,30 @@ def flash_create_args(self, extra_args=""):
         return args
 
     def flash_probe(self):
+        """
+        Execute flashrom with no commands to simply probe the flash chip
+        """
         args = self.flash_create_args()
         return self.flash_cmd(args)
 
     def flash_read(self, read_file):
+        """
+        Execute flashrom with read command to read the firmware from the DUT
+        """
         args = self.flash_create_args(f"-r {self.FW_PATH_READ}")
         return self.flash_cmd(args, read_file=read_file)
 
     def flash_erase(self):
+        """
+        Execute flashrom with erase command to erase the flash chip
+        """
         args = self.flash_create_args(f"-E")
         return self.flash_cmd(args)
 
     def flash_write(self, write_file, bios=False):
+        """
+        Execute flashrom with write command to write firmware to the DUT
+        """
         if "disable_wp" in self.dut_data:
             args = self.flash_create_args("--wp-disable --wp-range=0x0,0x0")
             self.flash_cmd(args)