From 65abd1006bca13aa5fc12b173bf206dc261bd592 Mon Sep 17 00:00:00 2001
From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Date: Fri, 16 Jul 2021 19:21:34 -0700
Subject: [PATCH] add context manager to temporarily change the cwd (#2377)

Commit history before merge:

* add context manager to temporarily change the cwd
* Iterator, not Iterable
---
 tests/test_black.py | 43 ++++++++++++++++---------------------------
 tests/util.py       | 14 +++++++++++++-
 2 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/tests/test_black.py b/tests/test_black.py
index e3be9c7..998ecfc 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -44,6 +44,7 @@ from pathspec import PathSpec
 # Import other test classes
 from tests.util import (
     THIS_DIR,
+    change_directory,
     read_data,
     DETERMINISTIC_HEADER,
     BlackBaseTestCase,
@@ -2009,17 +2010,12 @@ class BlackTestCase(BlackBaseTestCase):
             return
 
         # https://bugs.python.org/issue33660
-
-        old_cwd = Path.cwd()
-        try:
-            root = Path("/")
-            os.chdir(str(root))
+        root = Path("/")
+        with change_directory(root):
             path = Path("workspace") / "project"
             report = black.Report(verbose=True)
             normalized_path = black.normalize_path_maybe_ignore(path, root, report)
             self.assertEqual(normalized_path, "workspace/project")
-        finally:
-            os.chdir(str(old_cwd))
 
     def test_newline_comment_interaction(self) -> None:
         source = "class A:\\\r\n# type: ignore\n pass\n"
@@ -2170,10 +2166,6 @@ class BlackTestCase(BlackBaseTestCase):
         Test that the code option finds the pyproject.toml in the current directory.
         """
         with patch.object(black, "parse_pyproject_toml", return_value={}) as parse:
-            # Make sure we are in the project root with the pyproject file
-            if not Path("tests").exists():
-                os.chdir("..")
-
             args = ["--code", "print"]
             CliRunner().invoke(black.main, args)
 
@@ -2192,22 +2184,19 @@ class BlackTestCase(BlackBaseTestCase):
         Test that the code option finds the pyproject.toml in the parent directory.
         """
         with patch.object(black, "parse_pyproject_toml", return_value={}) as parse:
-            # Make sure we are in the tests directory
-            if Path("tests").exists():
-                os.chdir("tests")
-
-            args = ["--code", "print"]
-            CliRunner().invoke(black.main, args)
-
-            pyproject_path = Path(Path().cwd().parent, "pyproject.toml").resolve()
-            assert (
-                len(parse.mock_calls) >= 1
-            ), "Expected config parse to be called with the current directory."
-
-            _, call_args, _ = parse.mock_calls[0]
-            assert (
-                call_args[0].lower() == str(pyproject_path).lower()
-            ), "Incorrect config loaded."
+            with change_directory(Path("tests")):
+                args = ["--code", "print"]
+                CliRunner().invoke(black.main, args)
+
+                pyproject_path = Path(Path().cwd().parent, "pyproject.toml").resolve()
+                assert (
+                    len(parse.mock_calls) >= 1
+                ), "Expected config parse to be called with the current directory."
+
+                _, call_args, _ = parse.mock_calls[0]
+                assert (
+                    call_args[0].lower() == str(pyproject_path).lower()
+                ), "Incorrect config loaded."
 
 
 with open(black.__file__, "r", encoding="utf-8") as _bf:
diff --git a/tests/util.py b/tests/util.py
index 1e86a3f..e83017f 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -1,7 +1,8 @@
 import os
 import unittest
 from pathlib import Path
-from typing import List, Tuple, Any
+from typing import Iterator, List, Tuple, Any
+from contextlib import contextmanager
 from functools import partial
 
 import black
@@ -72,3 +73,14 @@ def read_data_from_file(file_name: Path) -> Tuple[str, str]:
         # If there's no output marker, treat the entire file as already pre-formatted.
         _output = _input[:]
     return "".join(_input).strip() + "\n", "".join(_output).strip() + "\n"
+
+
+@contextmanager
+def change_directory(path: Path) -> Iterator[None]:
+    """Context manager to temporarily chdir to a different directory."""
+    previous_dir = os.getcwd()
+    try:
+        os.chdir(path)
+        yield
+    finally:
+        os.chdir(previous_dir)
-- 
2.39.5