Fix #3539: reload module with assertion rewrite import hook

AUTHORS

@@ -93,6 +93,7 @@ Hui Wang (coldnight)
 Ian Bicking
 Ian Lesperance
 Ionuț Turturică
+Iwan Briquemont
 Jaap Broekhuizen
 Jan Balster
 Janne Vanhala

changelog/3539.bugfix.rst

@@ -0,0 +1 @@
+Fix reload on assertion rewritten modules.

src/_pytest/assertion/rewrite.py

@@ -269,13 +269,13 @@ class AssertionRewritingHook(object):
         )
 
     def load_module(self, name):
+        co, pyc = self.modules.pop(name)
+        if name in sys.modules:
             # If there is an existing module object named 'fullname' in
             # sys.modules, the loader must use that existing module. (Otherwise,
             # the reload() builtin will not work correctly.)
-        if name in sys.modules:
-            return sys.modules[name]
-
-        co, pyc = self.modules.pop(name)
+            mod = sys.modules[name]
+        else:
             # I wish I could just call imp.load_compiled here, but __file__ has to
             # be set properly. In Python 3.2+, this all would be handled correctly
             # by load_compiled.

testing/test_assertrewrite.py

@@ -1050,6 +1050,48 @@ class TestAssertionRewriteHookDetails(object):
         result = testdir.runpytest("-s")
         result.stdout.fnmatch_lines(["* 1 passed*"])
 
+    def test_reload_reloads(self, testdir):
+        """Reloading a module after change picks up the change."""
+        testdir.tmpdir.join("file.py").write(
+            textwrap.dedent(
+                """
+            def reloaded():
+                return False
+
+            def rewrite_self():
+                with open(__file__, 'w') as self:
+                    self.write('def reloaded(): return True')
+        """
+            )
+        )
+        testdir.tmpdir.join("pytest.ini").write(
+            textwrap.dedent(
+                """
+            [pytest]
+            python_files = *.py
+        """
+            )
+        )
+
+        testdir.makepyfile(
+            test_fun="""
+            import sys
+            try:
+                from imp import reload
+            except ImportError:
+                pass
+
+            def test_loader():
+                import file
+                assert not file.reloaded()
+                file.rewrite_self()
+                reload(file)
+                assert file.reloaded()
+            """
+        )
+        result = testdir.runpytest("-s")
+        result.stdout.fnmatch_lines(["* 1 passed*"])
+
     def test_get_data_support(self, testdir):
         """Implement optional PEP302 api (#808).
         """