| # Copyright 2022 The IREE Authors |
| # |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. |
| # See https://llvm.org/LICENSE.txt for license information. |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| from typing import Tuple |
| |
| from io import StringIO |
| import unittest |
| import sys |
| |
| from iree.runtime.scripts.iree_devices import __main__ as cli |
| |
| |
| def run_cli(*args) -> Tuple[int, str, str]: |
| capture_stdout = StringIO() |
| capture_stderr = StringIO() |
| sys.stdout = capture_stdout |
| sys.stderr = capture_stderr |
| try: |
| rc = cli.main(args) |
| finally: |
| sys.stdout = sys.__stdout__ |
| sys.stderr = sys.__stderr__ |
| return rc, capture_stdout.getvalue(), capture_stderr.getvalue() |
| |
| |
| class DevicesCliTest(unittest.TestCase): |
| def testLs(self): |
| rc, output, err = run_cli("ls") |
| self.assertEqual(rc, 0) |
| self.assertIn("vmvx:0\tdefault", output) |
| |
| def testLsTryCreate(self): |
| rc, output, err = run_cli("ls", "--try-create") |
| self.assertEqual(rc, 0) |
| self.assertIn("vmvx:0\tdefault\tSUCCESS", output) |
| |
| def testLsTryCreateExplicitDriver(self): |
| rc, output, err = run_cli("ls", "--try-create", "-d", "vmvx") |
| self.assertEqual(rc, 0) |
| self.assertIn("vmvx:0\tdefault\tSUCCESS", output) |
| |
| def testLsTryCreateExplicitDriverNotFound(self): |
| rc, output, err = run_cli("ls", "--try-create", "-d", "DOES_NOT_EXIST") |
| self.assertEqual(rc, 0) |
| self.assertIn("Could not create driver DOES_NOT_EXIST", err) |
| |
| def testTestIndexedDevice(self): |
| rc, output, err = run_cli("test", "vmvx:0") |
| self.assertEqual(rc, 0) |
| self.assertIn("Creating device vmvx:0... SUCCESS", output) |
| |
| def testTestDefaultDevice(self): |
| rc, output, err = run_cli("test", "vmvx") |
| self.assertEqual(rc, 0) |
| self.assertIn("Creating device vmvx... SUCCESS", output) |
| |
| def testTestNonExisting(self): |
| rc, output, err = run_cli("test", "NOT_EXISTING") |
| self.assertEqual(rc, 1) |
| self.assertIn("Creating device NOT_EXISTING... ERROR", output) |
| |
| |
| if __name__ == "__main__": |
| unittest.main() |
