cases: move main to avatar using a custom runner

5 files changed, 244 insertions, 43 deletions

diff --git a/.github/workflows/avatar.yml b/.github/workflows/avatar.yml
index f8dd9ba..9907d48 100644
--- a/.github/workflows/avatar.yml
+++ b/.github/workflows/avatar.yml
@@ -84,7 +84,7 @@ jobs:
         run: nohup python -m rootcanal > rootcanal.log &
       - name: Test
         run: |
-          python cases/main.py -l | grep -Ev '^=' > test-names.txt
-          timeout 5m python cases/main.py -c cases/config.yml --test_bed bumble.bumbles --tests $(split test-names.txt -n l/${{ matrix.shard }})
+          avatar cases/ --list | grep -Ev '^=' > test-names.txt
+          timeout 5m avatar cases/ --test-beds bumble.bumbles --tests $(split test-names.txt -n l/${{ matrix.shard }})
       - name: Rootcanal Logs
         run: cat rootcanal.log
diff --git a/avatar/__init__.py b/avatar/__init__.py
index aff4556..9141754 100644
--- a/avatar/__init__.py
+++ b/avatar/__init__.py
@@ -19,12 +19,14 @@ any Bluetooth test cases virtually and physically.
 
 __version__ = "0.0.2"
 
+import argparse
 import enum
 import functools
 import grpc
 import grpc.aio
 import importlib
 import logging
+import pathlib
 
 from avatar import pandora_server
 from avatar.aio import asynchronous
@@ -32,7 +34,9 @@ from avatar.metrics import trace
 from avatar.pandora_client import BumblePandoraClient as BumblePandoraDevice
 from avatar.pandora_client import PandoraClient as PandoraDevice
 from avatar.pandora_server import PandoraServer
+from avatar.runner import SuiteRunner
 from mobly import base_test
+from mobly import signals
 from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Sized, Tuple, Type, TypeVar
 
 # public symbols
@@ -106,7 +110,13 @@ class PandoraDevices(Sized, Iterable[PandoraDevice]):
 
             # Register the controller and load its Pandora servers.
             logging.info('Starting %s(s) for %s', server_cls.__name__, controller)
-            devices: Optional[List[Any]] = test.register_controller(server_cls.MOBLY_CONTROLLER_MODULE)  # type: ignore
+            try:
+                devices: Optional[List[Any]] = test.register_controller(  # type: ignore
+                    server_cls.MOBLY_CONTROLLER_MODULE
+                )
+            except Exception:
+                logging.exception('abort: failed to register controller')
+                raise signals.TestAbortAll("")
             assert devices
             for device in devices:  # type: ignore
                 self._servers.append(server_cls(device))
@@ -216,3 +226,87 @@ def rpc_except(
         return wrapper
 
     return wrap
+
+
+def args_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description='Avatar test runner.')
+    parser.add_argument(
+        'input',
+        type=str,
+        nargs='*',
+        metavar='<PATH>',
+        help='Lits of folder or test file to run',
+        default=['.'],
+    )
+    parser.add_argument('-c', '--config', type=str, metavar='<PATH>', help='Path to the test configuration file.')
+    parser.add_argument(
+        '-l',
+        '--list',
+        '--list_tests',  # For backward compatibility with tradefed `MoblyBinaryHostTest`
+        action='store_true',
+        help='Print the names of the tests defined in a script without ' 'executing them.',
+    )
+    parser.add_argument(
+        '-o',
+        '--log-path',
+        '--log_path',  # For backward compatibility with tradefed `MoblyBinaryHostTest`
+        type=str,
+        metavar='<PATH>',
+        help='Path to the test configuration file.',
+    )
+    parser.add_argument(
+        '-t',
+        '--tests',
+        nargs='+',
+        type=str,
+        metavar='[ClassA[.test_a] ClassB[.test_b] ...]',
+        help='A list of test classes and optional tests to execute.',
+    )
+    parser.add_argument(
+        '-b',
+        '--test-beds',
+        '--test_bed',  # For backward compatibility with tradefed `MoblyBinaryHostTest`
+        nargs='+',
+        type=str,
+        metavar='[<TEST BED NAME1> <TEST BED NAME2> ...]',
+        help='Specify which test beds to run tests on.',
+    )
+
+    parser.add_argument('-v', '--verbose', action='store_true', help='Set console logger level to DEBUG')
+    return parser
+
+
+# Avatar default entry point
+def main(args: Optional[argparse.Namespace] = None) -> None:
+    import sys
+
+    # Create an Avatar suite runner.
+    runner = SuiteRunner()
+
+    # Parse arguments.
+    argv = args or args_parser().parse_args()
+    if argv.input:
+        for path in argv.input:
+            runner.add_path(pathlib.Path(path))
+    if argv.config:
+        runner.add_config_file(pathlib.Path(argv.config))
+    if argv.log_path:
+        runner.set_logs_dir(pathlib.Path(argv.log_path))
+    if argv.tests:
+        runner.add_test_filters(argv.tests)
+    if argv.test_beds:
+        runner.add_test_beds(argv.test_beds)
+    if argv.verbose:
+        runner.set_logs_verbose()
+
+    # List tests to standard output.
+    if argv.list:
+        for _, (tag, test_names) in runner.included_tests.items():
+            for name in test_names:
+                print(f"{tag}.{name}")
+        sys.exit(0)
+
+    # Run the test suite.
+    logging.basicConfig(level=logging.INFO)
+    if not runner.run():
+        sys.exit(1)
diff --git a/avatar/runner.py b/avatar/runner.py
new file mode 100644
index 0000000..0b2a170
--- /dev/null
+++ b/avatar/runner.py
@@ -0,0 +1,141 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Avatar runner."""
+
+import inspect
+import logging
+import os
+import pathlib
+
+from importlib.machinery import SourceFileLoader
+from mobly import base_test
+from mobly import config_parser
+from mobly import signals
+from mobly import test_runner
+from typing import Dict, List, Tuple, Type
+
+_BUMBLE_BTSNOOP_FMT = 'bumble_btsnoop_{pid}_{instance}.log'
+
+
+class SuiteRunner:
+    test_beds: List[str] = []
+    test_run_configs: List[config_parser.TestRunConfig] = []
+    test_classes: List[Type[base_test.BaseTestClass]] = []
+    test_filters: List[str] = []
+    logs_dir: pathlib.Path
+    logs_verbose: bool = False
+
+    def __init__(self) -> None:
+        self.set_logs_dir(pathlib.Path('out'))
+
+    def set_logs_dir(self, path: pathlib.Path) -> None:
+        if not path.exists():
+            path.mkdir()
+        self.logs_dir = path
+
+    def set_logs_verbose(self, verbose: bool = True) -> None:
+        self.logs_verbose = verbose
+
+    def add_test_beds(self, test_beds: List[str]) -> None:
+        self.test_beds += test_beds
+
+    def add_test_filters(self, test_filters: List[str]) -> None:
+        self.test_filters += test_filters
+
+    def add_config_file(self, path: pathlib.Path) -> None:
+        self.test_run_configs += config_parser.load_test_config_file(str(path))  # type: ignore
+
+    def add_test_class(self, cls: Type[base_test.BaseTestClass]) -> None:
+        self.test_classes.append(cls)
+
+    def add_test_module(self, path: pathlib.Path) -> None:
+        try:
+            module = SourceFileLoader(path.stem, str(path)).load_module()
+            classes = inspect.getmembers(module, inspect.isclass)
+            for _, cls in classes:
+                if issubclass(cls, base_test.BaseTestClass):
+                    self.test_classes.append(cls)
+        except ImportError:
+            pass
+
+    def add_path(self, path: pathlib.Path, root: bool = True) -> None:
+        if path.is_file():
+            if path.name.endswith('_test.py'):
+                self.add_test_module(path)
+            elif not self.test_run_configs and not root and path.name in ('config.yml', 'config.yaml'):
+                self.add_config_file(path)
+            elif root:
+                raise ValueError(f'{path} is not a test file')
+        else:
+            for child in path.iterdir():
+                self.add_path(child, root=False)
+
+    def is_included(self, cls: base_test.BaseTestClass, test: str) -> bool:
+        return not self.test_filters or any(filter_match(cls, test, filter) for filter in self.test_filters)
+
+    @property
+    def included_tests(self) -> Dict[Type[base_test.BaseTestClass], Tuple[str, List[str]]]:
+        result: Dict[Type[base_test.BaseTestClass], Tuple[str, List[str]]] = {}
+        for test_class in self.test_classes:
+            cls = test_class(config_parser.TestRunConfig())
+            test_names: List[str] = []
+            try:
+                # Executes pre-setup procedures, this is required since it might
+                # generate test methods that we want to return as well.
+                cls._pre_run()
+                test_names = cls.tests or cls.get_existing_test_names()  # type: ignore
+                test_names = list(test for test in test_names if self.is_included(cls, test))
+                if test_names:
+                    assert cls.TAG
+                    result[test_class] = (cls.TAG, test_names)
+            except Exception:
+                logging.exception('Failed to retrieve generated tests.')
+            finally:
+                cls._clean_up()
+        return result
+
+    def run(self) -> bool:
+        # Enable Bumble snoop logs.
+        os.environ.setdefault('BUMBLE_SNOOPER', f'btsnoop:file:{self.logs_dir}/{_BUMBLE_BTSNOOP_FMT}')
+
+        # Execute the suite
+        ok = True
+        for config in self.test_run_configs:
+            test_bed: str = config.test_bed_name  # type: ignore
+            if self.test_beds and test_bed not in self.test_beds:
+                continue
+            runner = test_runner.TestRunner(config.log_path, config.testbed_name)
+            with runner.mobly_logger(console_level=logging.DEBUG if self.logs_verbose else logging.INFO):
+                for test_class, (_, tests) in self.included_tests.items():
+                    runner.add_test_class(config, test_class, tests)  # type: ignore
+                try:
+                    runner.run()
+                    ok = ok and runner.results.is_all_pass
+                except signals.TestAbortAll:
+                    ok = ok and not self.test_beds
+                except Exception:
+                    logging.exception('Exception when executing %s.', config.testbed_name)
+                    ok = False
+        return ok
+
+
+def filter_match(cls: base_test.BaseTestClass, test: str, filter: str) -> bool:
+    tag: str = cls.TAG  # type: ignore
+    if '.test_' in filter:
+        return f"{tag}.{test}".startswith(filter)
+    if filter.startswith('test_'):
+        return test.startswith(filter)
+    return tag.startswith(filter)
diff --git a/cases/main.py b/cases/main.py
deleted file mode 100644
index bbdb801..0000000
--- a/cases/main.py
+++ /dev/null
@@ -1,40 +0,0 @@
-import argparse
-import logging
-import os
-
-from argparse import Namespace
-from mobly import suite_runner
-from typing import List, Tuple
-
-_BUMBLE_BTSNOOP_FMT = 'bumble_btsnoop_{pid}_{instance}.log'
-
-# Import test cases modules.
-import host_test
-import le_host_test
-import le_security_test
-import security_test
-
-_TEST_CLASSES_LIST = [
-    host_test.HostTest,
-    le_host_test.LeHostTest,
-    security_test.SecurityTest,
-    le_security_test.LeSecurityTest,
-]
-
-
-def _parse_cli_args() -> Tuple[Namespace, List[str]]:
-    parser = argparse.ArgumentParser(description='Avatar test runner.')
-    parser.add_argument('-o', '--log_path', type=str, metavar='<PATH>', help='Path to the test configuration file.')
-    return parser.parse_known_args()
-
-
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
-
-    # Enable bumble snoop logger.
-    ns, argv = _parse_cli_args()
-    if ns.log_path:
-        os.environ.setdefault('BUMBLE_SNOOPER', f'btsnoop:file:{ns.log_path}/{_BUMBLE_BTSNOOP_FMT}')
-
-    # Run the test suite.
-    suite_runner.run_suite(_TEST_CLASSES_LIST, argv)  # type: ignore
diff --git a/pyproject.toml b/pyproject.toml
index 331b95e..a6079af 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,6 +20,9 @@ dependencies = [
 [project.urls]
 Source = "https://github.com/google/avatar"
 
+[project.scripts]
+avatar = "avatar:main"
+
 [project.optional-dependencies]
 dev = [
     "rootcanal==1.3.0",
@@ -35,6 +38,9 @@ dev = [
 [tool.flit.module]
 name = "avatar"
 
+[tool.flit.sdist]
+include = ["cases/", "doc/"]
+
 [tool.black]
 line-length = 119
 target-version = ["py38", "py39", "py310", "py311"]