Simplify audio sink iteration in docking station tests (New)

checkbox-support/checkbox_support/scripts/pipewire_utils.py

@@ -22,6 +22,7 @@
 import json
 import logging
 import re
+import shlex
 import subprocess
 import sys
 import time
@@ -78,7 +79,7 @@ class PipewireTest:
 
     logger = logging.getLogger()
 
-    def _get_pw_type(self, media_class) -> str:
+    def _get_pw_type(self, media_class: str) -> str:
         """
         convert sink to Output and source to Input
 
@@ -96,7 +97,9 @@ def _get_pw_type(self, media_class) -> str:
             self.logger.info("Media class:[{}] is unknown".format(media_class))
             return "UNKNOWN CLASS"
 
-    def _get_pw_dump(self, p_type) -> dict:
+    def _get_pw_dump(
+        self, p_type: 't.Literal["Device", "Node"]'
+    ) -> "list[dict[str, t.Any]]":
         """
         Use to convert the json output of pw-dump to dict object
 
@@ -113,7 +116,7 @@ def _get_pw_dump(self, p_type) -> dict:
             return json.loads(pw_dump)
         except (json.decoder.JSONDecodeError, TypeError):
             self.logger.error("pw-dump {} failed !!!".format(p_type))
-            return {}
+            return []
 
     def generate_pw_media_class(self, media_type, media_class) -> str:
         """
@@ -319,7 +322,7 @@ def gst_pipeline(self, pipe, timeout, device) -> int:
 
         return PipewireTestError.NO_ERROR
 
-    def _get_audio_config(self, mode) -> set:
+    def _get_audio_config(self, mode):
         """
         Get simple audio configuration
         This function parse output of pw-dump to find the device type
@@ -331,7 +334,7 @@ def _get_audio_config(self, mode) -> set:
         :type mode: str
         """
         clients = self._get_pw_dump("Device")
-        cfg = set()
+        cfg = set()  # type: set[tuple[str, str, str]]
         for client in clients:
             active_ports = None
             mclass = client["info"]["props"].get("media.class")
@@ -380,7 +383,7 @@ def monitor_active_port_change(self, timeout, mode) -> int:
         self.logger.info("Couldn't detect active port change!")
         return PipewireTestError.NO_CHANGE_DETECTED
 
-    def go_through_ports(self, cmd, mode):
+    def go_through_ports(self, cmd: str, mode: 't.Literal["source", "sink"]'):
         """
         Go through available ports for testing
         This script checks if the ports on either sinks
@@ -395,46 +398,182 @@ def go_through_ports(self, cmd, mode):
         """
         clients = self._get_pw_dump("Device")
         for client in clients:
-            ports = None
+            ports = []
             mclass = client["info"]["props"].get("media.class")
             if mclass == "Audio/Device":
                 ports = client["info"]["params"]["EnumRoute"]
-            if ports:
-                for p in ports:
-                    chosen = None
-                    if p["direction"] == self._get_pw_type(mode) and p[
-                        "available"
-                    ] in [
-                        "yes",
-                        "unknown",
-                    ]:
-                        while chosen != "yes":
-                            self.logger.info(
-                                "Please select [{}] for "
-                                "testing (if selected, "
-                                "please enter 'yes')".format(p["description"])
-                            )
+                assert type(ports) is list
+
+            for p in ports:
+                chosen = None
+                if p["direction"] == self._get_pw_type(mode) and p[
+                    "available"
+                ] in [
+                    "yes",
+                    "unknown",
+                ]:
+                    while chosen != "yes":
+                        self.logger.info(
+                            "Please select [{}] for "
+                            "testing (if selected, "
+                            "please enter 'yes')".format(p["description"])
+                        )
 
-                            chosen = input()
-                        checked = None
-                        while checked != "yes":
-                            with subprocess.Popen(
-                                cmd,
-                                shell=True,
-                                stdout=subprocess.PIPE,
-                                universal_newlines=True,
-                            ) as p:
-                                while p.poll() is None:
-                                    line = p.stdout.readline().strip()
-                                    self.logger.info(line)
-                                p.kill()
-                            self.logger.info(
-                                "Is working ?  please enter 'yes' to leave"
+                        chosen = input()
+                    checked = None
+                    while checked != "yes":
+                        # check_call will print to stdout for us
+                        subprocess.check_call(cmd, shell=True)
+                        self.logger.info(
+                            "Is working ?  please enter 'yes' to leave"
+                        )
+                        checked = input()
+
+    def iter_audio_sinks(self, cmd: "list[str]"):
+        """Execute the cmd for each audio sink discovered by pipewire
+
+        :param cmd: the command to run
+        """
+
+        tested_ids = set()  # type: set[int]
+        audio_sink_ids = list(self._find_available_audio_sinks().items())
+        N = len(audio_sink_ids)
+
+        if N == 0:
+            raise SystemExit("No audio sinks are available for this test")
+
+        while True:
+            try:
+                for i, (node_id, node_description) in enumerate(
+                    audio_sink_ids
+                ):
+                    print(
+                        "({}) - '{}' {}".format(
+                            i,
+                            node_description,
+                            ("- Tested" if node_id in tested_ids else ""),
+                        )
+                    )
+                _input = input(
+                    "Choose an audio sink to test [0-{}]".format(N - 1)
+                    + ", or type 'q' to quit: "
+                )
+                if _input == "q":
+                    if len(tested_ids) == N:
+                        print(
+                            "[ OK ] Quitting with return code 0.",
+                            "All {} audio sinks have been tested".format(N),
+                        )
+                        return
+                    else:
+                        raise SystemExit(
+                            "Only {} audio sinks were tested, ".format(
+                                len(tested_ids)
                             )
+                            + "but expected {}".format(N)
+                        )
+
+                idx = int(_input)
+                subprocess.check_call(
+                    ["wpctl", "set-default", str(audio_sink_ids[idx][0])]
+                )
+            except (ValueError, IndexError):
+                # this would loop at input() until a valid index is selected
+                print(
+                    "Please select an index from 0 to", N - 1, file=sys.stderr
+                )
+                continue
+
+            node_id, node_description = audio_sink_ids[idx]
 
-                            checked = input()
+            print("=" * 80, flush=True)
+            print(
+                "Testing '{}', id={}, command={}, 60s timeout".format(
+                    node_description, node_id, cmd
+                )
+            )
+            # don't let this fail, just go to the next sink
+            subprocess.run(cmd, timeout=60)
+            print("=" * 80, flush=True)
+
+            tested_ids.add(audio_sink_ids[idx][0])
+            print(
+                "Progress: {}/{} audio sinks tested".format(len(tested_ids), N)
+            )
 
-    def _get_node_description(self, properties) -> str:
+    def _find_available_audio_sinks(self) -> "dict[int, str]":
+        """
+        Finds the list of audio "devices" as shown in gnome's control center
+
+        :return: Returns a set of IDs that can be consumed by wpctl.
+                 The values are human readable names.
+                 These IDs are the "ID" to use as shown in `wpctl --help`
+        """
+        testable_node_ids = {}  # type: dict[int, str]
+        pw_audio_devices = [
+            device
+            for device in self._get_pw_dump("Device")
+            if device["info"]["props"].get("media.class") == "Audio/Device"
+        ]
+        pw_sink_nodes = [
+            node
+            for node in self._get_pw_dump("Node")
+            if node["info"]["props"].get("media.class") == "Audio/Sink"
+        ]
+
+        for node in pw_sink_nodes:
+            # IDs of these "nodes" can be passed to wpctl set-default
+            node_id = int(node["id"])
+            device_id = int(node["info"]["props"]["device.id"])
+
+            device = None  # type: dict[str, t.Any] | None
+            for dev in pw_audio_devices:
+                if dev["id"] == device_id:
+                    device = dev
+                    break
+
+            if not device:
+                print("Could not find device", device_id, file=sys.stderr)
+                continue
+
+            # now check if the device has at least 1 available route
+            enum_routes = device["info"]["params"]["EnumRoute"]
+
+            if not isinstance(enum_routes, list):
+                raise TypeError(
+                    "EnumRoute of device {} is not a list, got {}".format(
+                        device_id, type(enum_routes)
+                    )
+                )
+
+            for route in enum_routes:
+                # try to match the device to this node
+                if (
+                    route["devices"][0]  # this is an array with just 1 value
+                    != node["info"]["props"]["card.profile.device"]
+                ):
+                    continue
+                if route["direction"] != "Output":
+                    print(
+                        "Skipping '{}'".format(route["description"]),
+                        "because it's not a sink",
+                    )
+                    continue
+                if route["available"] not in ("yes", "unknown"):
+                    print(
+                        "Skipping '{}'".format(route["description"]),
+                        "because it's unavailable",
+                    )
+                    continue
+
+                # correct direction + at least 1 available route => testable
+                testable_node_ids[node_id] = "{} - {}".format(
+                    route["description"],
+                    device["info"]["props"].get("device.product.name", ""),
+                )
+        return testable_node_ids
+
+    def _get_node_description(self, properties) -> "str | None":
         """
         Get node description from the output of wpctl inspect
 
@@ -490,7 +629,7 @@ def show_default_device(self, device_type):
         except subprocess.CalledProcessError as e:
             raise RuntimeError("Show default device error {}".format(repr(e)))
 
-    def _sort_wpctl_status(self, lines: list) -> list:
+    def _sort_wpctl_status(self, lines: "list[str]") -> "list[str]":
         """
         This method will sort wpctl status for sub-items under catalog only
 
@@ -533,17 +672,17 @@ def compare_wpctl_status(self, status_1: str, status_2: str):
         :param status_2: path to second wpctl status
         """
         with open(status_1, "r") as s1, open(status_2, "r") as s2:
-            status_1 = s1.readlines()
-            status_2 = s2.readlines()
-            sorted_status_1 = self._sort_wpctl_status(status_1)
-            sorted_status_2 = self._sort_wpctl_status(status_2)
+            status_1_lines = s1.readlines()
+            status_2_lines = s2.readlines()
+            sorted_status_1 = self._sort_wpctl_status(status_1_lines)
+            sorted_status_2 = self._sort_wpctl_status(status_2_lines)
             delta = difflib.unified_diff(sorted_status_1, sorted_status_2, n=0)
             diff = "".join(delta)
             if diff:
                 self.logger.info("The first status:\n")
-                self.logger.info("".join(status_1))
+                self.logger.info("".join(status_1_lines))
                 self.logger.info("And the second status:\n")
-                self.logger.info("".join(status_2))
+                self.logger.info("".join(status_2_lines))
                 self.logger.info(
                     "Differ in the following lines (after sorting):"
                 )
@@ -753,6 +892,17 @@ def _args_parsing(self, args=sys.argv[1:]):
             "-m", "--mode", type=str, help="Either sinks or sources"
         )
 
+        parser_iter_sink = subparsers.add_parser(
+            "iter-audio-sinks", help="Iterate all available audio sinks"
+        )
+        parser_iter_sink.add_argument(
+            "-c",
+            "--command",
+            type=str,
+            required=True,
+            help="command for testing",
+        )
+
         # Add parser for show default device function
         parser_show = subparsers.add_parser(
             "show", help="show the default device"
@@ -823,6 +973,8 @@ def function_select(self, args):
                 return PipewireTestError.NO_ERROR
             else:
                 return PipewireTestError.NOT_REAL_DEVICE
+        elif args.test_type == "iter-audio-sinks":
+            return self.iter_audio_sinks(shlex.split(args.command))
 
 
 def main():