Add support for nested test case names (#12)

Author: pmeier

pytest_pytorch/plugin.py

@@ -1,28 +1,86 @@
+import re
 import unittest.mock
+from typing import Pattern
 
 from _pytest.unittest import TestCaseFunction, UnitTestCase
 
+from torch.testing._internal.common_device_type import get_device_type_test_bases
 from torch.testing._internal.common_utils import TestCase as PyTorchTestCaseTemplate
 
 
+class PytestPyTorchInternalError(Exception):
+    def __init__(self, msg):
+        super().__init__(
+            f"{msg}\n"
+            f"This is an internal error of the pytest plugin 'pytest-pytorch'."
+            f"If you encounter this during normal operation, please file an issue "
+            f"https://github.com/Quansight/pytest-pytorch/issues."
+        )
+
+
+TEST_BASE_DEVICE_PATTERN = re.compile(r"(?P<device>\w*?)TestBase$")
+
+
+def _get_devices():
+    devices = []
+    for test_base in get_device_type_test_bases():
+        match = TEST_BASE_DEVICE_PATTERN.match(test_base.__name__)
+        if not match:
+            raise PytestPyTorchInternalError(
+                f"Unable to extract device name from {test_base.__name__}"
+            )
+
+        devices.append(match.group("device"))
+
+    return devices
+
+
+DEVICES = _get_devices()
+
+
 class TemplateName(str):
+    _TEMPLATE_NAME_PATTERN: Pattern
+
+    def __init__(self, _):
+        super().__init__()
+        match = self._TEMPLATE_NAME_PATTERN.match(self)
+        if not match:
+            raise PytestPyTorchInternalError(
+                f"Unable to extract template name from {self}"
+            )
+        self._template_name = match.group("template_name")
+
     def __eq__(self, other):
-        return self.startswith(str(other))
+        return str.__eq__(self, other) or str.__eq__(self._template_name, other)
 
     def __hash__(self) -> int:
         return super().__hash__()
 
 
+class TestCaseFunctionTemplateName(TemplateName):
+    _TEMPLATE_NAME_PATTERN = re.compile(
+        fr"(?P<template_name>\w*?)_({'|'.join([device.lower() for device in DEVICES])})"
+    )
+
+
 class PyTorchTestCaseFunction(TestCaseFunction):
     @classmethod
     def from_parent(cls, parent, *, name, **kw):
-        return super().from_parent(parent, name=TemplateName(name), **kw)
+        return super().from_parent(
+            parent, name=TestCaseFunctionTemplateName(name), **kw
+        )
+
+
+class TestCaseTemplateName(TemplateName):
+    _TEMPLATE_NAME_PATTERN = re.compile(
+        fr"(?P<template_name>\w*?)({'|'.join([device.upper() for device in DEVICES])})"
+    )
 
 
 class PyTorchTestCase(UnitTestCase):
     @classmethod
     def from_parent(cls, parent, *, name, obj=None):
-        return super().from_parent(parent, name=TemplateName(name), obj=obj)
+        return super().from_parent(parent, name=TestCaseTemplateName(name), obj=obj)
 
     def collect(self):
         # Yes, this is a bad practice. Unfortunately, there is no other option to

tests/assets/test_nested_names.py

@@ -0,0 +1,24 @@
+from torch.testing._internal.common_device_type import (
+    instantiate_device_type_tests,
+    onlyOn,
+)
+from torch.testing._internal.common_utils import TestCase
+
+
+class TestFoo(TestCase):
+    # fails for meta, passes for cpu
+    def test_baz(self, device):
+        assert device != "meta"
+
+
+instantiate_device_type_tests(TestFoo, globals())
+
+
+class TestFooBar(TestCase):
+    # passes for meta, skips for cpu
+    @onlyOn("meta")
+    def test_baz(self, device):
+        assert True
+
+
+instantiate_device_type_tests(TestFooBar, globals())

tests/test_plugin.py

@@ -180,3 +180,34 @@ def test_dtype(testdir, file, cmds, outcomes):
     testdir.copy_example(file)
     result = testdir.runpytest(*cmds)
     result.assert_outcomes(**outcomes)
+
+
+@make_params(
+    "test_nested_names.py",
+    Config(
+        "*testcase-*test-*device",
+        new_cmds=(),
+        legacy_cmds=(),
+        passed=2,
+        skipped=1,
+        failed=1,
+    ),
+    Config(
+        "1testcase1-*test-*device",
+        new_cmds="::TestFoo",
+        legacy_cmds=("-k", "TestFoo and not TestFooBar"),
+        passed=1,
+        failed=1,
+    ),
+    Config(
+        "1testcase2-*test-*device",
+        new_cmds="::TestFooBar",
+        legacy_cmds=("-k", "TestFooBar"),
+        passed=1,
+        skipped=1,
+    ),
+)
+def test_nested_names(testdir, file, cmds, outcomes):
+    testdir.copy_example(file)
+    result = testdir.runpytest(*cmds)
+    result.assert_outcomes(**outcomes)