teststubtest: further fix tests

Basically a follow up to python#12282

mypy/test/teststubtest.py

@@ -24,7 +24,7 @@ def use_tmp_dir(mod_name: str) -> Iterator[None]:
             os.chdir(tmp)
             if sys.path[0] != tmp:
                 sys.path.insert(0, tmp)
-            yield
+            yield tmp
         finally:
             sys.path = current_syspath[:]
             if mod_name in sys.modules:
@@ -99,7 +99,7 @@ def staticmethod(f: T) -> T: ...
 def run_stubtest(
     stub: str, runtime: str, options: List[str], config_file: Optional[str] = None,
 ) -> str:
-    with use_tmp_dir(TEST_MODULE_NAME):
+    with use_tmp_dir(TEST_MODULE_NAME) as tmp_dir:
         with open("builtins.pyi", "w") as f:
             f.write(stubtest_builtins_stub)
         with open("typing.pyi", "w") as f:
@@ -118,10 +118,8 @@ def run_stubtest(
                 parse_options([TEST_MODULE_NAME] + options),
                 use_builtins_fixtures=True
             )
-
-        module_path = Path(os.getcwd()) / TEST_MODULE_NAME
         # remove cwd as it's not available from outside
-        return output.getvalue().replace(str(module_path), TEST_MODULE_NAME)
+        return output.getvalue().replace(tmp_dir + "/", "")
 
 
 class Case:
@@ -780,27 +778,18 @@ def test_dunders(self) -> Iterator[Case]:
                 error=None,
             )
 
-    def test_not_subclassable(self) -> None:
-        output = run_stubtest(
-            stub=(
-                "class CanBeSubclassed: ...\n"
-                "class CanNotBeSubclassed:\n"
-                "  def __init_subclass__(cls) -> None: ...\n"
-            ),
-            runtime=(
-                "class CanNotBeSubclassed:\n"
-                "  def __init_subclass__(cls): raise TypeError('nope')\n"
-                # ctypes.Array can be subclassed, but subclasses must define a few
-                # special attributes, e.g. _length_
-                "from ctypes import Array as CanBeSubclassed\n"
-            ),
-            options=[],
+    @collect_cases
+    def test_not_subclassable(self) -> Iterator[Case]:
+        yield Case(
+            stub="class CanBeSubclassed: ...",
+            runtime="class CanBeSubclassed: ...",
+            error=None,
+        )
+        yield Case(
+            stub="class CannotBeSubclassed:\n  def __init_subclass__(cls) -> None: ...",
+            runtime="class CannotBeSubclassed:\n  def __init_subclass__(cls): raise TypeError",
+            error="CannotBeSubclassed",
         )
-        assert (
-            "CanNotBeSubclassed cannot be subclassed at runtime,"
-            " but isn't marked with @final in the stub"
-        ) in output
-        assert "CanBeSubclassed cannot be subclassed" not in output
 
     @collect_cases
     def test_name_mangling(self) -> Iterator[Case]:
@@ -1114,7 +1103,7 @@ def test_config_file(self) -> None:
             "plugins={}/test-data/unit/plugins/decimal_to_int.py\n".format(root_dir)
         )
         output = run_stubtest(stub=stub, runtime=runtime, options=[])
-        assert output == (
+        assert remove_color_code(output) == (
             "error: test_module.temp variable differs from runtime type Literal[5]\n"
             "Stub: at line 2\ndecimal.Decimal\nRuntime:\n5\n\n"
         )