From ce70144718f0c3df5e64d183cbca7e9bd43a1c77 Mon Sep 17 00:00:00 2001 From: hanhsuan <32028620+hanhsuan@users.noreply.github.com> Date: Fri, 13 Sep 2024 17:53:57 +0800 Subject: [PATCH] Add a new function to find which renderer is running the process (New) (#1245) * While testing without pci id, user couldn't know the process is running on which renderer. Therefore, this change is going to: * add a new function to get which GPU is executing the command that could help user to make sure the renderer is not llvmpipe. * refactor some codes to make reuse easier. * add PCI ID in the log to make verify easier. * Fix typo * Update providers/base/bin/prime_offload_tester.py Co-authored-by: Fernando Bravo <39527354+fernando79513@users.noreply.github.com> * Update providers/base/bin/prime_offload_tester.py Co-authored-by: Fernando Bravo <39527354+fernando79513@users.noreply.github.com> * Update providers/base/bin/prime_offload_tester.py Co-authored-by: Fernando Bravo <39527354+fernando79513@users.noreply.github.com> * Update providers/base/bin/prime_offload_tester.py Co-authored-by: Fernando Bravo <39527354+fernando79513@users.noreply.github.com> * Update providers/base/bin/prime_offload_tester.py Co-authored-by: Fernando Bravo <39527354+fernando79513@users.noreply.github.com> * 1. using os.walk and open to replace grep and cat subprocess 2. removing _run_command 3. using run_with_timeout in the checkbox support to replace timeout command * remove useless code and fix black issue * Update providers/base/bin/prime_offload_tester.py Co-authored-by: Fernando Bravo <39527354+fernando79513@users.noreply.github.com> * 1. rename _if_file_contain 2. add description for cmd[0] * Fix black issue * Fixed unit tests * Mock test_nv_driver_check * Mocking correct timeout function * Add more unit tests * fix unit test error * Fix unit test error --------- Co-authored-by: Fernando Bravo <39527354+fernando79513@users.noreply.github.com> Co-authored-by: ferbraher --- providers/base/bin/prime_offload_tester.py | 298 ++++++---- .../base/tests/test_prime_offload_tester.py | 514 ++++++++++++------ 2 files changed, 550 insertions(+), 262 deletions(-) diff --git a/providers/base/bin/prime_offload_tester.py b/providers/base/bin/prime_offload_tester.py index 8fd7135832..7fc5a531fb 100755 --- a/providers/base/bin/prime_offload_tester.py +++ b/providers/base/bin/prime_offload_tester.py @@ -17,14 +17,16 @@ # You should have received a copy of the GNU General Public License # along with Checkbox. If not, see . -import sys -import threading +from checkbox_support.helpers.timeout import run_with_timeout import subprocess -import time -import re -import json +import threading import argparse import logging +import fnmatch +import time +import json +import sys +import re import os @@ -44,54 +46,77 @@ class PrimeOffloader: logger = logging.getLogger() check_result = False - def find_card_id(self, pci_name: str) -> str: + def find_file_containing_string( + self, search_directory: str, filename_pattern: str, search_string: str + ) -> str: """ - use pci name to find card id under /sys/kernel/debug/dri + Search for a file matching a specific pattern + that contains a given string. - :param pci_name: pci device name in NNNN:NN:NN.N format + :param search_directory: The directory to search through. + + :param filename_pattern: The pattern that filenames should match. + + :param search_string: The string to search for within + the file's contents. + + :returns: The full path of the file that contains the string, + or None if no file is found. + """ + for root, dirs, files in os.walk(search_directory): + for file_name in fnmatch.filter(files, filename_pattern): + file_path = os.path.join(root, file_name) + # Check if the search string is in the file + with open( + file_path, "r", encoding="utf-8", errors="ignore" + ) as file: + if search_string in file.read(): + return file_path + + def find_card_id(self, pci_bdf: str) -> str: + """ + use pci BDF to find card id under /sys/kernel/debug/dri + + :param pci_bdf: pci device BDF in NNNN:NN:NN.N format :returns: card id """ - pci_name_format = "[0-9]{4}:[0-9,a-f]{2}:[0-9,a-f]{2}.[0-9]" - if not re.match(pci_name_format, pci_name.lower()): - raise SystemExit("pci name format error") + pci_bdf_format = "[0-9]{4}:[0-9,a-f]{2}:[0-9,a-f]{2}.[0-9]" + if not re.match(pci_bdf_format, pci_bdf.lower()): + raise SystemExit("pci BDF format error") try: - cmd = [ - "grep", - "-lr", - "--include=name", - pci_name, - "/sys/kernel/debug/dri", - ] - - card_path = subprocess.check_output(cmd, universal_newlines=True) + card_path = self.find_file_containing_string( + "/sys/kernel/debug/dri", "name", pci_bdf + ) return card_path.split("/")[5] except IndexError as e: raise SystemExit("return value format error {}".format(repr(e))) - except subprocess.CalledProcessError as e: - raise SystemExit("run command failed {}".format(repr(e))) - def find_card_name(self, pci_name: str) -> str: + def find_card_name(self, pci_bdf: str) -> str: """ - use pci name to find card name by lshw + use pci BDF to find card name by lshw - :param pci_name: pci device name in NNNN:NN:NN.N format + :param pci_bdf: pci device BDF in NNNN:NN:NN.N format :returns: card name """ - cmd = ["lshw", "-c", "display", "-json"] + cmd = ["lshw", "-c", "display", "-numeric", "-json"] try: - card_infos = subprocess.check_output(cmd, universal_newlines=True) + card_infos = subprocess.check_output( + cmd, shell=False, universal_newlines=True + ) infos = json.loads(card_infos) for info in infos: - if pci_name in info["businfo"]: + if pci_bdf in info["businfo"]: return info["product"] raise SystemExit("Card name not found") except (KeyError, TypeError, json.decoder.JSONDecodeError) as e: raise SystemExit("return value format error {}".format(e)) - except subprocess.CalledProcessError as e: - raise SystemExit("run command failed {}".format(repr(e))) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + raise SystemExit( + "Running command:{} failed due to {}".format(cmd, repr(e)) + ) def get_clients(self, card_id: str) -> str: """ @@ -106,30 +131,26 @@ def get_clients(self, card_id: str) -> str: and the process could be found in /sys/kernel/debug/dri//clients - :param cmd: command that running under prime offload + :param card_id: card id shows in debugfs """ - read_clients_cmd = [ - "cat", - "/sys/kernel/debug/dri/{}/clients".format(card_id), - ] - try: - return subprocess.check_output( - read_clients_cmd, universal_newlines=True - ) - except subprocess.CalledProcessError: - self.logger.info( - "Couldn't get clients on specific GPU{}".format(card_id) - ) + filename = "/sys/kernel/debug/dri/{}/clients".format(card_id) + with open(filename, "r") as f: + return f.read() + return "" def check_offload( self, cmd: list, card_id: str, card_name: str, timeout: int ): """ - Check provided command is executed on specific GPU. - :param cmd: command to check if it's running on specific GPU - :param card_id: card id of dri device - :param card_name: card name of dri device - :param timeout: timeout for offloaded command + Used to check if the provided command is executed on a specific GPU. + + :param cmd: command to be run under prime offload + + :param card_id: card ID of the DRI device + + :param card_name: card name of the DRI device + + :param timeout: timeout for the offloaded command """ delay = timeout / 10 @@ -138,6 +159,8 @@ def check_offload( while time.time() < deadline: time.sleep(delay) clients = self.get_clients(card_id) + # The command shows in /sys/kernel/debug/dri//clients + # doesn't include arguments. Therefore cmd[0] is used to search if clients and cmd[0] in clients: self.logger.info("Checking success:") self.logger.info(" Offload process:[{}]".format(cmd)) @@ -148,6 +171,68 @@ def check_offload( self.logger.info(" Couldn't find process [{}]".format(cmd)) self.check_result = True + def _find_bdf(self, card_id: str): + """ + Use the /sys/kernel/debug/dri//name to get pci BDF. + + :param card_id: card id shows in debugfs + """ + filename = "/sys/kernel/debug/dri/{}/name".format(card_id) + with open(filename, "r") as f: + data_in_name = f.read() + return data_in_name.split()[1].split("=")[1] + + def find_offload(self, cmd: str, timeout: int): + """ + Find the card that the command is running on. + This script looks for the card on which a specific command is running. + It checks ten times in regular intervals if the process is running + by looking for it in the /sys/kernel/debug/dri//clients. + If the timeout is reached, the function will fail + + :param cmd: command that is running + + :param timeout: timeout for command + """ + directory = "/sys/kernel/debug/dri" + + delay = timeout / 10 + + deadline = time.time() + timeout + + cmd = cmd.split() + + while time.time() < deadline: + time.sleep(delay) + # The command shows in /sys/kernel/debug/dri//clients + # doesn't include arguments. Therefore cmd[0] is used to search + card_path = self.find_file_containing_string( + directory, "clients", cmd[0] + ) + if directory in card_path: + try: + # The graphic will be shown such as 0 and 128 + # at the same time. Therefore, pick up the first one + first_card = card_path.splitlines()[0] + card_id = first_card.split("/")[5] + bdf = self._find_bdf(card_id) + self.logger.info("Process is running on:") + self.logger.info(" process:[{}]".format(cmd)) + self.logger.info( + " Card ID:[{}]".format(self.find_card_id(bdf)) + ) + self.logger.info( + " Device Name:[{}]".format(self.find_card_name(bdf)) + ) + return + except IndexError as e: + self.logger.info( + "Finding card information failed {}".format(repr(e)) + ) + self.logger.info("Checking fail:") + self.logger.info(" Couldn't find process [{}]".format(cmd)) + self.check_result = True + def check_nv_offload_env(self): """ prime offload of nvidia driver is limited. @@ -177,44 +262,78 @@ def check_nv_offload_env(self): "No prime-select, it should be ok to run prime offload" ) - def run_offload_cmd( - self, cmd: str, pci_name: str, driver: str, timeout: int - ): + def cmd_runner(self, cmd: list, env: dict = None): + """ + use to execute command and piping the output to the screen. + + :param cmd: the command will be executed + + :param env: the environment variables for executing command + + """ + try: + with subprocess.Popen( + cmd, + env=env, + stdout=subprocess.PIPE, + universal_newlines=True, + ) as runner: + + self.logger.info("running command:[{}]".format(cmd)) + + # redirect command output real time + while runner.poll() is None: + line = runner.stdout.readline().strip() + self.logger.info(line) + except subprocess.CalledProcessError as e: + raise SystemExit("run command failed {}".format(repr(e))) + + def cmd_finder(self, cmd: str, timeout: int): + """ + run offload command and find it runs on which GPU + + :param cmd: command that running under prime offload + + :param timeout: timeout for offloaded command + """ + if "timeout" in cmd: + raise SystemExit("Put timeout in command isn't allowed") + + # use other thread to find offload + find_thread = threading.Thread( + target=self.find_offload, args=(cmd, timeout) + ) + find_thread.start() + try: + run_with_timeout(self.cmd_runner, timeout, cmd.split()) + except TimeoutError: + self.logger.info("Test finished") + find_thread.join() + + if self.check_result: + raise SystemExit("Couldn't find process running on GPU") + + def cmd_checker(self, cmd: str, pci_bdf: str, driver: str, timeout: int): """ run offload command and check it runs on correct GPU :param cmd: command that running under prime offload - :param pci_name: pci device name in NNNN:NN:NN.N format + :param pci_bdf: pci device name in NNNN:NN:NN.N format :param driver: GPU driver, such as i915, amdgpu, nvidia :param timeout: timeout for offloaded command """ - card_id = self.find_card_id(pci_name) - card_name = self.find_card_name(pci_name) + card_id = self.find_card_id(pci_bdf) + card_name = self.find_card_name(pci_bdf) # run offload command in other process - dri_pci_name_format = re.sub("[:.]", "_", pci_name) + dri_pci_bdf_format = re.sub("[:.]", "_", pci_bdf) if "timeout" in cmd: raise SystemExit("Put timeout in command isn't allowed") - cmd = cmd.split() - if timeout > 0: - offload_cmd = ["timeout", str(timeout)] + cmd - else: - # if timeout <=0 will make check_offload failed. - # Set the timeout to the default value - log_str = ( - "Timeout {}s is invalid," - " remove the timeout setting" - " and change check_offload to run 20s".format(timeout) - ) - self.logger.info(log_str) - timeout = 20 - offload_cmd = cmd - env = os.environ.copy() if driver in ("nvidia", "pcieport"): offload_env = { @@ -222,7 +341,7 @@ def run_offload_cmd( "__GLX_VENDOR_LIBRARY_NAME": "nvidia", } else: - offload_env = {"DRI_PRIME": "pci-{}".format(dri_pci_name_format)} + offload_env = {"DRI_PRIME": "pci-{}".format(dri_pci_bdf_format)} env.update(offload_env) self.logger.info("prime offload env: {}".format(offload_env)) @@ -236,24 +355,13 @@ def run_offload_cmd( ) check_thread.start() try: - with subprocess.Popen( - offload_cmd, - env=env, - stdout=subprocess.PIPE, - universal_newlines=True, - ) as offload: - - self.logger.info("offload command:[{}]".format(offload_cmd)) + run_with_timeout(self.cmd_runner, timeout, cmd.split(), env) + except TimeoutError: + self.logger.info("Test finished") + check_thread.join() - # redirect offload command output real time - while offload.poll() is None: - line = offload.stdout.readline().strip() - self.logger.info(line) - check_thread.join() - if self.check_result: - raise SystemExit("offload to specific GPU failed") - except subprocess.CalledProcessError as e: - raise SystemExit("run offload command failed {}".format(repr(e))) + if self.check_result: + raise SystemExit("offload to specific GPU failed") def parse_args(self, args=sys.argv[1:]): """ @@ -278,15 +386,13 @@ def parse_args(self, args=sys.argv[1:]): "-p", "--pci", type=str, - default="0000:00:02.0", - help="pci name in NNNN:NN:NN.N format (default: %(default)s)", + help="pci device bdf in NNNN:NN:NN.N format, such as 0000:00:02.0", ) parser.add_argument( "-d", "--driver", type=str, - default="i915", - help="Type of GPU driver (default: %(default)s)", + help="Type of GPU driver, such as i915", ) parser.add_argument( "-t", @@ -313,8 +419,12 @@ def main(self): # Add console handler to logger self.logger.addHandler(console_handler) - # run_offload_cmd("glxgears", "0000:00:02.0", "i915", 0) - self.run_offload_cmd(args.command, args.pci, args.driver, args.timeout) + if args.pci and args.driver: + # cmd_checker("glxgears", "0000:00:02.0", "i915", 0) + self.cmd_checker(args.command, args.pci, args.driver, args.timeout) + else: + # cmd_finder("glxgears", 0) + self.cmd_finder(args.command, args.timeout) if __name__ == "__main__": diff --git a/providers/base/tests/test_prime_offload_tester.py b/providers/base/tests/test_prime_offload_tester.py index 733ffeedc8..1b0b23a37c 100755 --- a/providers/base/tests/test_prime_offload_tester.py +++ b/providers/base/tests/test_prime_offload_tester.py @@ -15,10 +15,45 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . +from unittest.mock import patch, MagicMock, mock_open +import subprocess import unittest -from unittest.mock import patch, MagicMock +import os -from prime_offload_tester import * +from checkbox_support.helpers.timeout import mock_timeout + +from prime_offload_tester import PrimeOffloader + + +class FindFileContainingString(unittest.TestCase): + """ + This function should work like "grep -l" + """ + + @patch("os.walk") + def test_grep_true(self, mock_walk): + mock_walk.return_value = [("/foo/bar", ["bar"], ["spam"])] + po = PrimeOffloader() + with patch("builtins.open", mock_open(read_data="test")) as mock_file: + self.assertEqual( + po.find_file_containing_string("/foo/bar", "spam", "test"), + "/foo/bar/spam", + ) + mock_file.assert_called_with( + "/foo/bar/spam", "r", encoding="utf-8", errors="ignore" + ) + + @patch("os.walk") + def test_grep_false(self, mock_walk): + mock_walk.return_value = [("/foo/bar", ["bar"], ["spam"])] + po = PrimeOffloader() + with patch("builtins.open", mock_open(read_data="test")) as mock_file: + self.assertEqual( + po.find_file_containing_string("/foo/bar", "spam", "xxx"), None + ) + mock_file.assert_called_with( + "/foo/bar/spam", "r", encoding="utf-8", errors="ignore" + ) class FindCardIdTests(unittest.TestCase): @@ -27,64 +62,65 @@ class FindCardIdTests(unittest.TestCase): (pci bus information) """ - @patch("subprocess.check_output") - def test_pci_name_format_check(self, mock_check): + @patch("prime_offload_tester.PrimeOffloader.find_file_containing_string") + def test_pci_name_format_check(self, mock_cmd): po = PrimeOffloader() # correct format - mock_check.return_value = "/sys/kernel/debug/dri/0/name" + mock_cmd.return_value = "/sys/kernel/debug/dri/0/name" self.assertEqual(po.find_card_id("0000:00:00.0"), "0") - mock_check.assert_called_with( - [ - "grep", - "-lr", - "--include=name", - "0000:00:00.0", - "/sys/kernel/debug/dri", - ], - universal_newlines=True, + mock_cmd.assert_called_with( + "/sys/kernel/debug/dri", + "name", + "0000:00:00.0", ) + + @patch("prime_offload_tester.PrimeOffloader.find_file_containing_string") + def test_pci_name_hex_format_check(self, mock_cmd): + po = PrimeOffloader() # should work with hex vaule + mock_cmd.return_value = "/sys/kernel/debug/dri/0/name" self.assertEqual(po.find_card_id("0000:c6:F0.0"), "0") + @patch("prime_offload_tester.PrimeOffloader.find_file_containing_string") + def test_pci_name_non_hex_format_check(self, mock_cmd): + po = PrimeOffloader() # error format - with alphabet + mock_cmd.return_value = "/sys/kernel/debug/dri/0/name" with self.assertRaises(SystemExit): po.find_card_id("000r:00:00.0") + @patch("prime_offload_tester.PrimeOffloader.find_file_containing_string") + def test_pci_name_digital_error_format_check(self, mock_cmd): + po = PrimeOffloader() # error format - digital position error + mock_cmd.return_value = "/sys/kernel/debug/dri/0/name" with self.assertRaises(SystemExit): po.find_card_id("0000:00:000.0") - @patch("subprocess.check_output") - def test_id_not_found(self, mock_check): + @patch("prime_offload_tester.PrimeOffloader.find_file_containing_string") + def test_empty_string_id_not_found(self, mock_cmd): po = PrimeOffloader() # empty string - mock_check.return_value = "" + mock_cmd.return_value = "" with self.assertRaises(SystemExit): po.find_card_id("0000:00:00.0") - mock_check.assert_called_with( - [ - "grep", - "-lr", - "--include=name", - "0000:00:00.0", - "/sys/kernel/debug/dri", - ], - universal_newlines=True, + mock_cmd.assert_called_with( + "/sys/kernel/debug/dri", + "name", + "0000:00:00.0", ) - # subprocess error - mock_check.side_effect = subprocess.CalledProcessError(-1, "test") + @patch("prime_offload_tester.PrimeOffloader.find_file_containing_string") + def test_indexerror_id_not_found(self, mock_cmd): + po = PrimeOffloader() + # IndexError + mock_cmd.side_effect = IndexError with self.assertRaises(SystemExit): po.find_card_id("0000:00:00.0") - mock_check.assert_called_with( - [ - "grep", - "-lr", - "--include=name", - "0000:00:00.0", - "/sys/kernel/debug/dri", - ], - universal_newlines=True, + mock_cmd.assert_called_with( + "/sys/kernel/debug/dri", + "name", + "0000:00:00.0", ) @@ -162,74 +198,107 @@ class FindCardNameTests(unittest.TestCase): ]""" @patch("subprocess.check_output") - def test_name_found_check(self, mock_check): + def test_name_found_check(self, mock_cmd): po = PrimeOffloader() - mock_check.return_value = self.lshw_output + mock_cmd.return_value = self.lshw_output self.assertEqual( po.find_card_name("0000:00:02.0"), "TigerLake-LP GT2 [Iris Xe Graphics]", ) - mock_check.assert_called_with( - ["lshw", "-c", "display", "-json"], universal_newlines=True + mock_cmd.assert_called_with( + ["lshw", "-c", "display", "-numeric", "-json"], + shell=False, + universal_newlines=True, ) @patch("subprocess.check_output") - def test_name_not_found_check(self, mock_check): + def test_pci_bdf_error_name_not_found_check(self, mock_cmd): po = PrimeOffloader() - # pci_name error - mock_check.return_value = self.lshw_output + # pci_bdf error + mock_cmd.return_value = self.lshw_output with self.assertRaises(SystemExit): po.find_card_name("0000:00:00.0") - mock_check.assert_called_with( - ["lshw", "-c", "display", "-json"], universal_newlines=True + mock_cmd.assert_called_with( + ["lshw", "-c", "display", "-numeric", "-json"], + shell=False, + universal_newlines=True, ) + @patch("subprocess.check_output") + def test_no_lshw_output_name_not_found_check(self, mock_cmd): + po = PrimeOffloader() # no businfo in lshw output - mock_check.return_value = self.lshw_output_err + mock_cmd.return_value = self.lshw_output_err with self.assertRaises(SystemExit): po.find_card_name("0000:00:00.0") - mock_check.assert_called_with( - ["lshw", "-c", "display", "-json"], universal_newlines=True + mock_cmd.assert_called_with( + ["lshw", "-c", "display", "-numeric", "-json"], + shell=False, + universal_newlines=True, ) + @patch("subprocess.check_output") + def test_empty_string_name_not_found_check(self, mock_cmd): + po = PrimeOffloader() # empty string - mock_check.return_value = "" + mock_cmd.return_value = "" with self.assertRaises(SystemExit): po.find_card_name("0000:00:00.0") - mock_check.assert_called_with( - ["lshw", "-c", "display", "-json"], universal_newlines=True + mock_cmd.assert_called_with( + ["lshw", "-c", "display", "-numeric", "-json"], + shell=False, + universal_newlines=True, ) + @patch("subprocess.check_output") + def test_none_name_not_found_check(self, mock_cmd): + po = PrimeOffloader() # None - mock_check.return_value = None + mock_cmd.return_value = None with self.assertRaises(SystemExit): po.find_card_name("0000:00:00.0") - mock_check.assert_called_with( - ["lshw", "-c", "display", "-json"], universal_newlines=True + mock_cmd.assert_called_with( + ["lshw", "-c", "display", "-numeric", "-json"], + shell=False, + universal_newlines=True, ) - # subprocess error - mock_check.side_effect = subprocess.CalledProcessError(-1, "test") + @patch("subprocess.check_output") + def test_keyerror_name_not_found_check(self, mock_cmd): + po = PrimeOffloader() + mock_cmd.side_effect = KeyError with self.assertRaises(SystemExit): po.find_card_name("0000:00:00.0") - mock_check.assert_called_with( - ["lshw", "-c", "display", "-json"], universal_newlines=True + mock_cmd.assert_called_with( + ["lshw", "-c", "display", "-numeric", "-json"], + shell=False, + universal_newlines=True, ) @patch("subprocess.check_output") - def test_get_clients(self, mock_check): + def test_subprocess_run_fail(self, mock_cmd): po = PrimeOffloader() - mock_check.return_value = "echo" - self.assertEqual(po.get_clients(0), "echo") - mock_check.assert_called_with( - ["cat", "/sys/kernel/debug/dri/0/clients"], universal_newlines=True + mock_cmd.side_effect = subprocess.CalledProcessError + with self.assertRaises(SystemExit): + po.find_card_name("0000:00:00.0") + mock_cmd.assert_called_with( + ["lshw", "-c", "display", "-numeric", "-json"], + shell=False, + universal_newlines=True, ) - # subprocess failed - mock_check.side_effect = subprocess.CalledProcessError(-1, "fail") - with self.assertRaises(subprocess.CalledProcessError): - po.check_nv_offload_env() - self.assertEqual(po.get_clients(0), None) + +class GetClientsTests(unittest.TestCase): + """ + This function should return the data in the right + "/sys/kernel/debug/dri/*/clients" + """ + + def test_get_clients(self): + po = PrimeOffloader() + with patch("builtins.open", mock_open(read_data="data")) as mock_file: + self.assertEqual(po.get_clients(0), "data") + mock_file.assert_called_with("/sys/kernel/debug/dri/0/clients", "r") class CheckOffloadTests(unittest.TestCase): @@ -249,13 +318,14 @@ def test_offload_succ_check(self, mock_client, mock_sleep): ) self.assertEqual(po.check_result, False) - @patch("time.time", side_effect=[1, 2, 3, 4]) @patch("time.sleep", return_value=None) + @patch("time.time") @patch("prime_offload_tester.PrimeOffloader.get_clients") - def test_offload_fail_check(self, mock_client, mock_sleep, mock_time): + def test_offload_fail_check(self, mock_client, mock_time, mock_sleep): cmd = ["echo"] # get_clients return string that doesn't include cmd mock_client.return_value = "" + mock_time.side_effect = [0, 0, 1, 2] po = PrimeOffloader() self.assertEqual( po.check_offload(cmd, "card_id", "card_name", 1), None @@ -264,6 +334,7 @@ def test_offload_fail_check(self, mock_client, mock_sleep, mock_time): # get_clients return None by CalledProcessError mock_client.return_value = None + mock_time.side_effect = [0, 0, 1, 2] po = PrimeOffloader() self.assertEqual( po.check_offload(cmd, "card_id", "card_name", 1), None @@ -271,6 +342,60 @@ def test_offload_fail_check(self, mock_client, mock_sleep, mock_time): self.assertEqual(po.check_result, True) +class FindBDFTests(unittest.TestCase): + """ + This function should return the BDF in the right + "/sys/kernel/debug/dri/*/name" + """ + + def test_find_bdf(self): + po = PrimeOffloader() + with patch( + "builtins.open", + mock_open(read_data="i915 dev=0000:00:02.0 unique=0000:00:02.0"), + ) as mock_file: + self.assertEqual(po._find_bdf(0), "0000:00:02.0") + mock_file.assert_called_with("/sys/kernel/debug/dri/0/name", "r") + + +class FindOffloadTests(unittest.TestCase): + """ + This function should try to find which GPU is the renderer + for the specific command. + """ + + @patch("time.sleep", return_value=None) + def test_found(self, mock_sleep): + cmd = "echo" + po = PrimeOffloader() + po._find_bdf = MagicMock(return_value="0") + po.find_card_id = MagicMock(return_value="0") + po.find_card_name = MagicMock(return_value="Intel") + po.find_file_containing_string = MagicMock( + return_value="/sys/kernel/debug/dri/0/clients" + ) + self.assertEqual(po.find_offload(cmd, 1), None) + self.assertEqual(po.check_result, False) + po.find_file_containing_string("/sys/kernel/debug/dri", "clients", cmd) + self.assertEqual(po.check_result, False) + + @patch("time.sleep", return_value=None) + @patch("time.time") + def test_not_found(self, mock_time, mock_sleep): + cmd = "echo" + po = PrimeOffloader() + po._find_bdf = MagicMock(return_value="0") + po.find_card_id = MagicMock(return_value="0") + po.find_card_name = MagicMock(return_value="Intel") + po.find_file_containing_string = MagicMock(return_value="") + mock_time.side_effect = [0, 0, 1, 2] + self.assertEqual(po.find_offload(cmd, 1), None) + po.find_file_containing_string.assert_called_with( + "/sys/kernel/debug/dri", "clients", cmd + ) + self.assertEqual(po.check_result, True) + + class CheckNvOffloadEnvTests(unittest.TestCase): """ This function will check this system could use prime offload or not. @@ -319,47 +444,131 @@ def test_nvlink_check(self, mock_check): self.assertEqual(po.check_nv_offload_env(), None) -class RunOffloadCmdTests(unittest.TestCase): +class CmdRunnerTests(unittest.TestCase): + """ + This function should run the command with Popen mode + and the right environment variables. + """ + + o_env = {"DRI_PRIME": "pci-0000_00_00_0"} + + @patch("subprocess.Popen") + def test_cmd_runner_succ(self, mock_open): + po = PrimeOffloader() + os.environ.copy = MagicMock(return_value={}) + po.cmd_runner(["echo", "0000:00:00.0", "xxx", 0], self.o_env) + # check cmd_runner executing correct command + mock_open.assert_called_with( + ["echo", "0000:00:00.0", "xxx", 0], + env=self.o_env, + stdout=subprocess.PIPE, + universal_newlines=True, + ) + + @patch("subprocess.Popen") + def test_cmd_runner_fail(self, mock_open): + po = PrimeOffloader() + mock_open.side_effect = subprocess.CalledProcessError(-1, "test") + with self.assertRaises(SystemExit): + po.cmd_runner(["echo", "0000:00:00.0", "xxx", 0], self.o_env) + + +@patch("subprocess.Popen", MagicMock()) +class CmdFinderTests(unittest.TestCase): + """ + This function should find the command is rendered by which GPU + """ + + def test_command_include_timeout(self): + po = PrimeOffloader() + po.find_offload = MagicMock(return_value="") + os.environ.copy = MagicMock(return_value={}) + po.check_result = True + with self.assertRaises(SystemExit): + po.cmd_finder("timeout 20 glxgears", 20) + + def test_found(self): + po = PrimeOffloader() + po.find_offload = MagicMock(return_value="") + os.environ.copy = MagicMock(return_value={}) + po.check_result = False + po.cmd_finder("glxgears", 20) + # check check_offload function get correct args + po.find_offload.assert_called_with("glxgears", 20) + + def test_not_found(self): + po = PrimeOffloader() + po.find_offload = MagicMock(return_value="") + os.environ.copy = MagicMock(return_value={}) + po.check_result = True + with self.assertRaises(SystemExit): + po.cmd_finder("glxgears", 20) + # check check_offload function get correct args + po.find_offload.assert_called_with("glxgears", 20) + + @patch("prime_offload_tester.run_with_timeout") + @patch("threading.Thread") + def test_not_found(self, mock_thread, mock_run_timeout): + po = PrimeOffloader() + po.check_offload = MagicMock(return_value="") + po.check_result = True + mock_run_timeout.side_effect = TimeoutError + with self.assertRaises(SystemExit): + po.cmd_finder("glxgears", 1) + # check check_offload function get correct args + mock_thread.assert_called_with( + target=po.find_offload, args=("glxgears", 1) + ) + + +@patch("subprocess.check_output", MagicMock()) +class CmdCheckerTests(unittest.TestCase): """ This function is the entry point to run the command with prime offload, if the environment is supported. """ - def test_condition_check(self): + nv_env = { + "__NV_PRIME_RENDER_OFFLOAD": "1", + "__GLX_VENDOR_LIBRARY_NAME": "nvidia", + } + o_env = {"DRI_PRIME": "pci-0000_00_00_0"} + + def test_no_card_id_check(self): po = PrimeOffloader() # no card id po.find_card_id = MagicMock(side_effect=SystemExit) with self.assertRaises(SystemExit): - po.run_offload_cmd("echo", "0000:00:00.0", "driver", 0) + po.cmd_checker("echo", "0000:00:00.0", "driver", 0) + def test_no_card_name_check(self): + po = PrimeOffloader() # no card name po.find_card_id = MagicMock(return_value="0") po.find_card_name = MagicMock(side_effect=SystemExit) with self.assertRaises(SystemExit): - po.run_offload_cmd("echo", "0000:00:00.0", "driver", 0) + po.cmd_checker("echo", "0000:00:00.0", "driver", 0) + def test_timeout_in_cmd_check(self): + po = PrimeOffloader() # timeout in command po.find_card_id = MagicMock(return_value="0") po.find_card_name = MagicMock(return_value="Card") with self.assertRaises(SystemExit): - po.run_offload_cmd("timeout 10 echo", "0000:00:00.0", "driver", 0) + po.cmd_checker("timeout 10 echo", "0000:00:00.0", "driver", 0) + @patch("prime_offload_tester.run_with_timeout", MagicMock()) + def test_nv_env_fail_check(self): + po = PrimeOffloader() # check_nv_offload_env failed po.find_card_id = MagicMock(return_value="0") po.find_card_name = MagicMock(return_value="Card") po.check_nv_offload_env = MagicMock(side_effect=SystemExit) with self.assertRaises(SystemExit): - po.run_offload_cmd("echo", "0000:00:00.0", "driver", 0) - - @patch("time.sleep", return_value=None) - @patch("subprocess.Popen") - def test_offload_cmd_check(self, mock_open, mock_sleep): - nv_env = { - "__NV_PRIME_RENDER_OFFLOAD": "1", - "__GLX_VENDOR_LIBRARY_NAME": "nvidia", - } - o_env = {"DRI_PRIME": "pci-0000_00_00_0"} + po.cmd_checker("echo", "0000:00:00.0", "driver", 0) + @patch("prime_offload_tester.run_with_timeout", MagicMock()) + def test_non_nv_driver_check(self): # non NV driver po = PrimeOffloader() po.find_card_id = MagicMock(return_value="0") @@ -367,89 +576,40 @@ def test_offload_cmd_check(self, mock_open, mock_sleep): po.check_nv_offload_env = MagicMock(return_value=None) po.check_offload = MagicMock(return_value="") os.environ.copy = MagicMock(return_value={}) - po.run_offload_cmd("echo", "0000:00:00.0", "xxx", 0) - # check run_offload_cmd executing correct command - mock_open.assert_called_with( - ["echo"], - env=o_env, - stdout=subprocess.PIPE, - universal_newlines=True, - ) - # check check_offload function get correct args - po.check_offload.assert_called_with(["echo"], "0", "Intel", 20) - - # non NV driver with timeout setting - po.find_card_id = MagicMock(return_value="0") - po.find_card_name = MagicMock(return_value="Intel") - po.check_nv_offload_env = MagicMock(return_value=None) - po.check_offload = MagicMock(return_value="") - os.environ.copy = MagicMock(return_value={}) - po.run_offload_cmd("echo", "0000:00:00.0", "xxx", 1) - # check run_offload_cmd executing correct command - mock_open.assert_called_with( - ["timeout", "1", "echo"], - env=o_env, - stdout=subprocess.PIPE, - universal_newlines=True, - ) + po.cmd_checker("glxgears", "0000:00:00.0", "xxx", 0) # check check_offload function get correct args - po.check_offload.assert_called_with(["echo"], "0", "Intel", 1) + po.check_offload.assert_called_with("glxgears", "0", "Intel", 0) + @patch("prime_offload_tester.run_with_timeout", MagicMock()) + def test_nv_driver_check(self): # NV driver + po = PrimeOffloader() po.find_card_id = MagicMock(return_value="0") po.find_card_name = MagicMock(return_value="NV") po.check_nv_offload_env = MagicMock(return_value=None) po.check_offload = MagicMock(return_value="") os.environ.copy = MagicMock(return_value={}) - po.run_offload_cmd("echo", "0000:00:00.0", "nvidia", 1) - # check run_offload_cmd executing correct command - mock_open.assert_called_with( - ["timeout", "1", "echo"], - env=nv_env, - stdout=subprocess.PIPE, - universal_newlines=True, - ) + po.cmd_checker("glxgears", "0000:00:00.0", "nvidia", 1) # check check_offload function get correct args - po.check_offload.assert_called_with(["echo"], "0", "NV", 1) + po.check_offload.assert_called_with("glxgears", "0", "NV", 1) - # subprocess error + @patch("prime_offload_tester.run_with_timeout") + @patch("threading.Thread") + def test_not_found(self, mock_thread, mock_run_timeout): + po = PrimeOffloader() po.find_card_id = MagicMock(return_value="0") po.find_card_name = MagicMock(return_value="NV") - po.check_nv_offload_env = MagicMock(return_value=None) po.check_offload = MagicMock(return_value="") - os.environ.copy = MagicMock(return_value={}) - mock_open.side_effect = subprocess.CalledProcessError(-1, "test") - with self.assertRaises(SystemExit): - po.run_offload_cmd("echo", "0000:00:00.0", "nvidia", 1) - # check run_offload_cmd executing correct command - mock_open.assert_called_with( - ["timeout", "1", "echo"], - env=nv_env, - stdout=subprocess.PIPE, - universal_newlines=True, - ) - # check check_offload function get correct args - po.check_offload.assert_called_with(["echo"], "0", "NV", 1) - - # check offload failed - po.find_card_id = MagicMock(return_value="0") - po.find_card_name = MagicMock(return_value="NV") po.check_nv_offload_env = MagicMock(return_value=None) - po.check_offload = MagicMock(return_value="") os.environ.copy = MagicMock(return_value={}) po.check_result = True - mock_open.side_effect = None + mock_run_timeout.side_effect = TimeoutError with self.assertRaises(SystemExit): - po.run_offload_cmd("echo", "0000:00:00.0", "nvidia", 1) - # check run_offload_cmd executing correct command - mock_open.assert_called_with( - ["timeout", "1", "echo"], - env=nv_env, - stdout=subprocess.PIPE, - universal_newlines=True, - ) + po.cmd_checker("glxgears", "0000:00:00.0", "nvidia", 1) # check check_offload function get correct args - po.check_offload.assert_called_with(["echo"], "0", "NV", 1) + mock_thread.assert_called_with( + target=po.check_offload, args=("glxgears", "0", "NV", 1) + ) class ParseArgsTests(unittest.TestCase): @@ -459,16 +619,16 @@ def test_success(self): args = [] rv = po.parse_args(args) self.assertEqual(rv.command, "glxgears") - self.assertEqual(rv.pci, "0000:00:02.0") - self.assertEqual(rv.driver, "i915") + self.assertEqual(rv.pci, None) + self.assertEqual(rv.driver, None) self.assertEqual(rv.timeout, 20) # change command args = ["-c", "glxgears -fullscreen"] rv = po.parse_args(args) self.assertEqual(rv.command, "glxgears -fullscreen") - self.assertEqual(rv.pci, "0000:00:02.0") - self.assertEqual(rv.driver, "i915") + self.assertEqual(rv.pci, None) + self.assertEqual(rv.driver, None) self.assertEqual(rv.timeout, 20) # change pci @@ -476,14 +636,14 @@ def test_success(self): rv = po.parse_args(args) self.assertEqual(rv.command, "glxgears") self.assertEqual(rv.pci, "0000:00:01.0") - self.assertEqual(rv.driver, "i915") + self.assertEqual(rv.driver, None) self.assertEqual(rv.timeout, 20) # change driver args = ["-d", "nvidia"] rv = po.parse_args(args) self.assertEqual(rv.command, "glxgears") - self.assertEqual(rv.pci, "0000:00:02.0") + self.assertEqual(rv.pci, None) self.assertEqual(rv.driver, "nvidia") self.assertEqual(rv.timeout, 20) @@ -491,8 +651,8 @@ def test_success(self): args = ["-t", "5"] rv = po.parse_args(args) self.assertEqual(rv.command, "glxgears") - self.assertEqual(rv.pci, "0000:00:02.0") - self.assertEqual(rv.driver, "i915") + self.assertEqual(rv.pci, None) + self.assertEqual(rv.driver, None) self.assertEqual(rv.timeout, 5) # change all @@ -515,18 +675,36 @@ def test_success(self): class MainTests(unittest.TestCase): @patch("prime_offload_tester.PrimeOffloader.parse_args") - @patch("prime_offload_tester.PrimeOffloader.run_offload_cmd") - def test_run_offload_cmd_succ(self, mock_run_offload, mock_parse_args): + @patch("prime_offload_tester.PrimeOffloader.cmd_checker") + def test_run_cmd_checker_succ(self, mock_checker, mock_parse_args): + args_mock = MagicMock() + args_mock.command = "cmd" + args_mock.pci = "pci" + args_mock.driver = "driver" + args_mock.timeout = "1" + mock_parse_args.return_value = args_mock self.assertEqual(PrimeOffloader().main(), None) + mock_checker.assert_called_with( + "cmd", + "pci", + "driver", + "1", + ) @patch("prime_offload_tester.PrimeOffloader.parse_args") - @patch("prime_offload_tester.PrimeOffloader.run_offload_cmd") - def test_run_offload_cmd_fail(self, mock_run_offload, mock_parse_args): - po = PrimeOffloader() - mock_run_offload.side_effect = SystemExit - with self.assertRaises(SystemExit) as cm: - po.main() - self.assertNotEqual(cm.exception.code, 0) + @patch("prime_offload_tester.PrimeOffloader.cmd_finder") + def test_run_cmd_finder_succ(self, mock_finder, mock_parse_args): + args_mock = MagicMock() + args_mock.command = "cmd" + args_mock.pci = None + args_mock.driver = None + args_mock.timeout = "1" + mock_parse_args.return_value = args_mock + self.assertEqual(PrimeOffloader().main(), None) + mock_finder.assert_called_with( + "cmd", + "1", + ) if __name__ == "__main__":