remote/client: Provide a way to set the strategy end-state

It is possible to set the initial state of a strategy, but not the final
state. In many cases we want to power the board off and otherwise clean
things up at the end. Provide a -e option to accomplish this.

Signed-off-by: Simon Glass <sjg@chromium.org>

Author: sjg20

labgrid/remote/client.py

@@ -884,6 +884,12 @@ def set_initial_state(self, target):
             logging.info("Transitioning into state %s", self.args.state)
             strategy.transition(self.args.state)
 
+    def set_end_state(self, target):
+        if self.args.end_state:
+            strategy = target.get_driver("Strategy")
+            logging.info("Transitioning into state %s", self.args.end_state)
+            strategy.transition(self.args.end_state)
+
     def _get_target(self, place):
         self._prepare_manager()
         target = None
@@ -1882,6 +1888,7 @@ def get_parser(auto_doc_mode=False) -> "argparse.ArgumentParser | AutoProgramArg
         type=str,
         help="strategy state to force into before switching to desired state",
     )
+    parser.add_argument("-e", "--end-state", type=str, help="strategy state to switch into after command")
     parser.add_argument(
         "-d", "--debug", action="store_true", default=False, help="enable debug mode (show python tracebacks)"
     )
@@ -2274,6 +2281,7 @@ def main():
     place = os.environ.get("LG_PLACE", place)
     state = os.environ.get("STATE", None)
     state = os.environ.get("LG_STATE", state)
+    end_state = os.environ.get("LG_END_STATE", state)
     initial_state = os.environ.get("LG_INITIAL_STATE", None)
     token = os.environ.get("LG_TOKEN", None)
 
@@ -2299,6 +2307,9 @@ def main():
     if args.initial_state is None:
         args.initial_state = initial_state
 
+    if args.end_state is None:
+        args.end_state = end_state
+
     if args.command in ["cancel-reservation", "wait"] and args.token is None:
         if token:
             args.token = token
@@ -2370,6 +2381,7 @@ def main():
             try:
                 if inspect.iscoroutinefunction(args.func):
                     auto_release = False
+                    target = None
                     if getattr(args.func, "needs_target", False):
                         if args.acquire:
                             place = session.get_place(args.place)
@@ -2384,6 +2396,7 @@ def main():
                     else:
                         coro = args.func(session)
                     session.loop.run_until_complete(coro)
+                    session.set_end_state(target)
                     if auto_release:
                         coro = session.release()
                         session.loop.run_until_complete(coro)

man/labgrid-client.1

@@ -40,7 +40,7 @@ This is the client to control a boards status and interface with it on remote ma
 .sp
 .EX
 usage: labgrid\-client [\-x ADDRESS] [\-c CONFIG] [\-p PLACE] [\-s STATE] [\-a]
-                      [\-i INITIAL_STATE] [\-d] [\-v] [\-P PROXY]
+                      [\-i INITIAL_STATE] [\-e END_STATE] [\-d] [\-v] [\-P PROXY]
                       COMMAND ...
 .EE
 .UNINDENT
@@ -77,6 +77,11 @@ strategy state to force into before switching to desired state
 .UNINDENT
 .INDENT 0.0
 .TP
+.B \-e <end_state>, \-\-end\-state <end_state>
+strategy state to switch into after command
+.UNINDENT
+.INDENT 0.0
+.TP
 .B \-d, \-\-debug
 enable debug mode (show python tracebacks)
 .UNINDENT