map module fullname to origin

Author: P403n1x87

src/_pytest/assertion/rewrite.py

@@ -80,7 +80,7 @@ def __init__(self, config: Config) -> None:
         self._basenames_to_check_rewrite = {"conftest"}
         self._marked_for_rewrite_cache: dict[str, bool] = {}
         self._session_paths_checked = False
-        self.fn: str | None = None
+        self._fns: dict[str, str] = {}
 
     def set_session(self, session: Session | None) -> None:
         self.session = session
@@ -127,7 +127,7 @@ def find_spec(
         ):
             return None
         else:
-            self.fn = fn = spec.origin
+            self._fns[name] = fn = spec.origin
 
         if not self._should_rewrite(name, fn, state):
             return None
@@ -145,8 +145,8 @@ def create_module(
         return None  # default behaviour is fine
 
     def get_code(self, fullname: str) -> types.CodeType:
-        assert self.fn is not None
-        fn = Path(self.fn)
+        assert fullname in self._fns
+        fn = Path(self._fns[fullname])
         state = self.config.stash[assertstate_key]
 
         # The requested module looks like a test file, so rewrite it. This is
@@ -185,11 +185,15 @@ def get_code(self, fullname: str) -> types.CodeType:
         return co
 
     def exec_module(self, module: types.ModuleType) -> None:
-        assert self.fn is not None
-
         module_name = module.__name__
 
-        self._rewritten_names[module_name] = Path(self.fn)
+        assert (
+            module_name in self._fns
+            and module.__spec__ is not None
+            and module.__spec__.origin == self._fns[module_name]
+        )
+
+        self._rewritten_names[module_name] = Path(self._fns[module_name])
 
         exec(self.get_code(module_name), module.__dict__)