Add more test cases for #506

tests/fixers/test_versioned_branches.py

@@ -8,6 +8,9 @@
 settings = Settings(target_version=(4, 0))
 check_noop = partial(tools.check_noop, settings=settings)
 check_transformed = partial(tools.check_transformed, settings=settings)
+check_error_on_transformed = partial(
+    tools.check_error_on_transformed, settings=settings
+)
 
 
 def test_future_version_gt():
@@ -263,3 +266,93 @@ def test_removed_block_internal_comment():
 
         """,
     )
+
+
+def test_removed_block_internal_comment1():
+    check_transformed(
+        """\
+        import django
+
+        if django.VERSION < (3, 2):
+            foo()
+            # test comment 1
+        # test comment 2
+        """,
+        """\
+        import django
+
+        # test comment 2
+        """,
+    )
+
+
+def test_removed_block_internal_comment2():
+    check_transformed(
+        """\
+        import django
+
+        # test comment 0
+        if django.VERSION < (3, 2):
+            foo()
+            # test comment 1
+        # test comment 2
+        foo()
+        """,
+        """\
+        import django
+
+        # test comment 0
+        # test comment 2
+        foo()
+        """,
+    )
+
+
+def test_removed_block_internal_comment3():
+    check_transformed(
+        """\
+        import django
+
+        if True:
+            # test comment 0
+            if django.VERSION < (3, 2):
+                foo()
+                # test comment 1
+            foo()
+        # test comment 2
+        foo()
+        """,
+        """\
+        import django
+
+        if True:
+            # test comment 0
+            foo()
+        # test comment 2
+        foo()
+        """,
+    )
+
+
+def test_removed_block_internal_comment_with_error():
+    check_error_on_transformed(
+        """\
+        import django
+
+        if True:
+            # test comment 0
+            if django.VERSION < (3, 2):
+                foo()
+                # test comment 1
+        # test comment 2
+        foo()
+        """,
+        """\
+        import django
+
+        if True:
+            # test comment 0
+        # test comment 2
+        foo()
+        """,
+    )

tests/fixers/tools.py

@@ -21,3 +21,15 @@ def check_transformed(
     ast.parse(dedented_after)  # check that the target is valid python code
     fixed = apply_fixers(dedented_before, settings=settings, filename=filename)
     assert fixed == dedented_after
+
+
+def check_error_on_transformed(
+    before: str, after: str, settings: Settings, filename: str = "example.py"
+) -> None:
+    dedented_after = dedent(after)
+    error = None
+    try:
+        ast.parse(dedented_after)  # check that the target is valid python code
+    except SyntaxError as ex:
+        error = ex
+    assert error is not None