fix test

deshaw · Dec 5, 2023 · dbd1a2f · dbd1a2f
1 parent a0d4bd7
commit dbd1a2f
Showing 1 changed file with 17 additions and 17 deletions.
diff --git a/tests/test_cmdline.py b/tests/test_cmdline.py
@@ -13,7 +13,7 @@
 import tempfile
 from   textwrap                 import dedent
 
-from   pyflyby._util            import EnvVarCtx
+from pyflyby._util import EnvVarCtx, CwdCtx
 
 import pytest
 
@@ -790,23 +790,23 @@ class B:
                 from foo import A, B
             """).lstrip())
             dot_pyflyby_fp.flush()
+        with CwdCtx(temp_dir):
+            result = pipe(
+                [BIN_DIR + "/tidy-imports", foo_fp.name],
+                env={"PYFLYBY_PATH": dot_pyflyby},
+            )
+
+            expected = dedent(
+                """
+                from __future__ import annotations
 
-        os.chdir(temp_dir)
-        result = pipe([
-            BIN_DIR+"/tidy-imports", foo_fp.name
-        ], env={
-            "PYFLYBY_PATH": dot_pyflyby
-        })
-
-        expected = dedent("""
-            from __future__ import annotations
-
-            class A:
-                param1: str
-                param2: B
+                class A:
+                    param1: str
+                    param2: B
 
 
-            class B:
-                param1: str
-        """).strip()
+                class B:
+                    param1: str
+            """
+            ).strip()
         assert result == expected