Add tests for switch_working_directory helper

tests/test_helpers.py

@@ -2,8 +2,17 @@
 
 import os
 
+from typing import TYPE_CHECKING
+
+import pytest
+
 from tests.helpers import flatten_dict
 from tests.helpers import isolated_environment
+from tests.helpers import switch_working_directory
+
+
+if TYPE_CHECKING:
+    from pathlib import Path
 
 
 def test_flatten_dict() -> None:
@@ -44,3 +53,29 @@ def test_isolated_environment_updates_environ() -> None:
     with isolated_environment(environ={"NEW_VAR": "new_value"}):
         assert os.environ["NEW_VAR"] == "new_value"
     assert "NEW_VAR" not in os.environ
+
+
+@pytest.mark.parametrize(
+    ("remove", "raise_error"),
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_switch_working_directory_changes_restores_and_removes(
+    tmp_path: Path, remove: bool, raise_error: bool
+) -> None:
+    original_cwd = os.getcwd()
+    temp_dir = tmp_path / f"temp-working-dir-{remove}-{raise_error}"
+    temp_dir.mkdir()
+
+    if raise_error:
+        with (
+            pytest.raises(RuntimeError),
+            switch_working_directory(temp_dir, remove=remove),
+        ):
+            assert os.getcwd() == str(temp_dir)
+            raise RuntimeError("boom")
+    else:
+        with switch_working_directory(temp_dir, remove=remove):
+            assert os.getcwd() == str(temp_dir)
+
+    assert os.getcwd() == original_cwd
+    assert temp_dir.exists() is (not remove)