Added test for exceptions in RemoveUnusedImports

pixee · Oct 26, 2023 · 73f5b4a · 73f5b4a
1 parent 39edc0f
commit 73f5b4a
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/tests/codemods/test_remove_unused_imports.py b/tests/codemods/test_remove_unused_imports.py
@@ -120,3 +120,15 @@ def test_ignore_init_files(self, tmpdir):
         tmp_file_path = Path(tmpdir / "__init__.py")
         self.run_and_assert_filepath(tmpdir, tmp_file_path, before, before)
         assert len(self.file_context.codemod_changes) == 0
+
+    def test_no_pyling_pragma_in_comment_trailing(self, tmpdir):
+        before = "import a # bogus: no-pragma"
+        after = ""
+        self.run_and_assert(tmpdir, before, after)
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_no_pyling_pragma_in_comment_before(self, tmpdir):
+        before = "#header\nprint('hello')\n# bogus: no-pragma\nimport a "
+        after = "#header\nprint('hello')"
+        self.run_and_assert(tmpdir, before, after)
+        assert len(self.file_context.codemod_changes) == 1