[Backend Tester] Clean up operator test logic (pytorch#12736)

Minor refactoring on the operator test logic - since we now have
separate model tests, I've moved operator-test specific test helper
logic into the operators directory and updated usages. I also updated
discovery slightly to give nicer error messages when things don't
import.

Author: GregoryComer

backends/test/suite/__init__.py

@@ -9,18 +9,11 @@
 
 import logging
 import os
-import unittest
-
-from enum import Enum
-from typing import Callable
 
 import executorch.backends.test.suite.flow
 
-import torch
-from executorch.backends.test.suite.context import get_active_test_context, TestContext
 from executorch.backends.test.suite.flow import TestFlow
-from executorch.backends.test.suite.reporting import log_test_summary
-from executorch.backends.test.suite.runner import run_test, runner_main
+from executorch.backends.test.suite.runner import runner_main
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -62,109 +55,6 @@ def get_test_flows() -> dict[str, TestFlow]:
     return _ALL_TEST_FLOWS
 
 
-DTYPES = [
-    # torch.int8,
-    # torch.uint8,
-    # torch.int16,
-    # torch.uint16,
-    # torch.int32,
-    # torch.uint32,
-    # torch.int64,
-    # torch.uint64,
-    # torch.float16,
-    torch.float32,
-    # torch.float64,
-]
-
-FLOAT_DTYPES = [
-    torch.float16,
-    torch.float32,
-    torch.float64,
-]
-
-
-# The type of test function. This controls the test generation and expected signature.
-# Standard tests are run, as is. Dtype tests get a variant generated for each dtype and
-# take an additional dtype parameter.
-class TestType(Enum):
-    STANDARD = 1
-    DTYPE = 2
-
-
-# Function annotation for dtype tests. This instructs the test framework to run the test
-# for each supported dtype and to pass dtype as a test parameter.
-def dtype_test(func):
-    func.test_type = TestType.DTYPE
-    return func
-
-
-# Class annotation for operator tests. This triggers the test framework to register
-# the tests.
-def operator_test(cls):
-    _create_tests(cls)
-    return cls
-
-
-# Generate test cases for each backend flow.
-def _create_tests(cls):
-    for key in dir(cls):
-        if key.startswith("test_"):
-            _expand_test(cls, key)
-
-
-# Expand a test into variants for each registered flow.
-def _expand_test(cls, test_name: str):
-    test_func = getattr(cls, test_name)
-    for flow in get_test_flows().values():
-        _create_test_for_backend(cls, test_func, flow)
-    delattr(cls, test_name)
-
-
-def _make_wrapped_test(
-    test_func: Callable,
-    test_name: str,
-    flow: TestFlow,
-    params: dict | None = None,
-):
-    def wrapped_test(self):
-        with TestContext(test_name, flow.name, params):
-            test_kwargs = params or {}
-            test_kwargs["flow"] = flow
-
-            test_func(self, **test_kwargs)
-
-    wrapped_test._name = test_name
-    wrapped_test._flow = flow
-
-    return wrapped_test
-
-
-def _create_test_for_backend(
-    cls,
-    test_func: Callable,
-    flow: TestFlow,
-):
-    test_type = getattr(test_func, "test_type", TestType.STANDARD)
-
-    if test_type == TestType.STANDARD:
-        wrapped_test = _make_wrapped_test(test_func, test_func.__name__, flow)
-        test_name = f"{test_func.__name__}_{flow.name}"
-        setattr(cls, test_name, wrapped_test)
-    elif test_type == TestType.DTYPE:
-        for dtype in DTYPES:
-            wrapped_test = _make_wrapped_test(
-                test_func,
-                test_func.__name__,
-                flow,
-                {"dtype": dtype},
-            )
-            dtype_name = str(dtype)[6:]  # strip "torch."
-            test_name = f"{test_func.__name__}_{dtype_name}_{flow.name}"
-            setattr(cls, test_name, wrapped_test)
-    else:
-        raise NotImplementedError(f"Unknown test type {test_type}.")
-
-
 def load_tests(loader, suite, pattern):
     package_dir = os.path.dirname(__file__)
     discovered_suite = loader.discover(
@@ -174,32 +64,5 @@ def load_tests(loader, suite, pattern):
     return suite
 
 
-class OperatorTest(unittest.TestCase):
-    def _test_op(self, model, inputs, flow: TestFlow):
-        context = get_active_test_context()
-
-        # This should be set in the wrapped test. See _make_wrapped_test above.
-        assert context is not None, "Missing test context."
-
-        run_summary = run_test(
-            model,
-            inputs,
-            flow,
-            context.test_name,
-            context.params,
-        )
-
-        log_test_summary(run_summary)
-
-        if not run_summary.result.is_success():
-            if run_summary.result.is_backend_failure():
-                raise RuntimeError("Test failure.") from run_summary.error
-            else:
-                # Non-backend failure indicates a bad test. Mark as skipped.
-                raise unittest.SkipTest(
-                    f"Test failed for reasons other than backend failure. Error: {run_summary.error}"
-                )
-
-
 if __name__ == "__main__":
     runner_main()

backends/test/suite/discovery.py

@@ -69,8 +69,17 @@ def _filter_tests(
 def _is_test_enabled(test_case: unittest.TestCase, test_filter: TestFilter) -> bool:
     test_method = getattr(test_case, test_case._testMethodName)
 
+    # Handle import / discovery failures - leave them enabled to report nicely at the
+    # top level. There might be a better way to do this. Internally, unittest seems to
+    # replace it with a stub method to report the failure.
+    if "testFailure" in str(test_method):
+        print(f"Warning: Test {test_case._testMethodName} failed to import.")
+        return True
+
     if not hasattr(test_method, "_flow"):
-        print(f"Test missing flow: {test_method}")
+        raise RuntimeError(
+            f"Test missing flow: {test_case._testMethodName} {test_method}"
+        )
 
     flow: TestFlow = test_method._flow
 

backends/test/suite/operators/__init__.py

@@ -7,6 +7,17 @@
 # pyre-unsafe
 
 import os
+import unittest
+
+from enum import Enum
+from typing import Callable
+
+import torch
+from executorch.backends.test.suite import get_test_flows
+from executorch.backends.test.suite.context import get_active_test_context, TestContext
+from executorch.backends.test.suite.flow import TestFlow
+from executorch.backends.test.suite.reporting import log_test_summary
+from executorch.backends.test.suite.runner import run_test
 
 
 def load_tests(loader, suite, pattern):
@@ -16,3 +27,133 @@ def load_tests(loader, suite, pattern):
     )
     suite.addTests(discovered_suite)
     return suite
+
+
+DTYPES = [
+    # torch.int8,
+    # torch.uint8,
+    # torch.int16,
+    # torch.uint16,
+    # torch.int32,
+    # torch.uint32,
+    # torch.int64,
+    # torch.uint64,
+    # torch.float16,
+    torch.float32,
+    # torch.float64,
+]
+
+FLOAT_DTYPES = [
+    torch.float16,
+    torch.float32,
+    torch.float64,
+]
+
+
+# The type of test function. This controls the test generation and expected signature.
+# Standard tests are run, as is. Dtype tests get a variant generated for each dtype and
+# take an additional dtype parameter.
+class TestType(Enum):
+    STANDARD = 1
+    DTYPE = 2
+
+
+# Function annotation for dtype tests. This instructs the test framework to run the test
+# for each supported dtype and to pass dtype as a test parameter.
+def dtype_test(func):
+    func.test_type = TestType.DTYPE
+    return func
+
+
+# Class annotation for operator tests. This triggers the test framework to register
+# the tests.
+def operator_test(cls):
+    _create_tests(cls)
+    return cls
+
+
+# Generate test cases for each backend flow.
+def _create_tests(cls):
+    for key in dir(cls):
+        if key.startswith("test_"):
+            _expand_test(cls, key)
+
+
+# Expand a test into variants for each registered flow.
+def _expand_test(cls, test_name: str):
+    test_func = getattr(cls, test_name)
+    for flow in get_test_flows().values():
+        _create_test_for_backend(cls, test_func, flow)
+    delattr(cls, test_name)
+
+
+def _make_wrapped_test(
+    test_func: Callable,
+    test_name: str,
+    flow: TestFlow,
+    params: dict | None = None,
+):
+    def wrapped_test(self):
+        with TestContext(test_name, flow.name, params):
+            test_kwargs = params or {}
+            test_kwargs["flow"] = flow
+
+            test_func(self, **test_kwargs)
+
+    wrapped_test._name = test_name
+    wrapped_test._flow = flow
+
+    return wrapped_test
+
+
+def _create_test_for_backend(
+    cls,
+    test_func: Callable,
+    flow: TestFlow,
+):
+    test_type = getattr(test_func, "test_type", TestType.STANDARD)
+
+    if test_type == TestType.STANDARD:
+        wrapped_test = _make_wrapped_test(test_func, test_func.__name__, flow)
+        test_name = f"{test_func.__name__}_{flow.name}"
+        setattr(cls, test_name, wrapped_test)
+    elif test_type == TestType.DTYPE:
+        for dtype in DTYPES:
+            wrapped_test = _make_wrapped_test(
+                test_func,
+                test_func.__name__,
+                flow,
+                {"dtype": dtype},
+            )
+            dtype_name = str(dtype)[6:]  # strip "torch."
+            test_name = f"{test_func.__name__}_{dtype_name}_{flow.name}"
+            setattr(cls, test_name, wrapped_test)
+    else:
+        raise NotImplementedError(f"Unknown test type {test_type}.")
+
+
+class OperatorTest(unittest.TestCase):
+    def _test_op(self, model, inputs, flow: TestFlow):
+        context = get_active_test_context()
+
+        # This should be set in the wrapped test. See _make_wrapped_test above.
+        assert context is not None, "Missing test context."
+
+        run_summary = run_test(
+            model,
+            inputs,
+            flow,
+            context.test_name,
+            context.params,
+        )
+
+        log_test_summary(run_summary)
+
+        if not run_summary.result.is_success():
+            if run_summary.result.is_backend_failure():
+                raise RuntimeError("Test failure.") from run_summary.error
+            else:
+                # Non-backend failure indicates a bad test. Mark as skipped.
+                raise unittest.SkipTest(
+                    f"Test failed for reasons other than backend failure. Error: {run_summary.error}"
+                )

backends/test/suite/operators/test_add.py

@@ -8,10 +8,14 @@
 
 
 import torch
-
-from executorch.backends.test.suite import dtype_test, operator_test, OperatorTest
 from executorch.backends.test.suite.flow import TestFlow
 
+from executorch.backends.test.suite.operators import (
+    dtype_test,
+    operator_test,
+    OperatorTest,
+)
+
 
 class Model(torch.nn.Module):
     def forward(self, x, y):

backends/test/suite/operators/test_div.py

@@ -10,10 +10,14 @@
 from typing import Optional
 
 import torch
-
-from executorch.backends.test.suite import dtype_test, operator_test, OperatorTest
 from executorch.backends.test.suite.flow import TestFlow
 
+from executorch.backends.test.suite.operators import (
+    dtype_test,
+    operator_test,
+    OperatorTest,
+)
+
 
 class Model(torch.nn.Module):
     def forward(self, x, y):

backends/test/suite/operators/test_elu.py

@@ -8,10 +8,14 @@
 
 
 import torch
-
-from executorch.backends.test.suite import dtype_test, operator_test, OperatorTest
 from executorch.backends.test.suite.flow import TestFlow
 
+from executorch.backends.test.suite.operators import (
+    dtype_test,
+    operator_test,
+    OperatorTest,
+)
+
 
 class Model(torch.nn.Module):
     def __init__(self, alpha=1.0, inplace=False):