diff --git a/crates/ruff/src/jupyter/notebook.rs b/crates/ruff/src/jupyter/notebook.rs
index 9e6a626f0aab4..5f75c60cd39c8 100644
--- a/crates/ruff/src/jupyter/notebook.rs
+++ b/crates/ruff/src/jupyter/notebook.rs
@@ -541,7 +541,7 @@ print("after empty cells")
     #[test]
     fn test_import_sorting() -> Result<()> {
         let path = "isort.ipynb".to_string();
-        let (diagnostics, source_kind) = test_notebook_path(
+        let (diagnostics, source_kind, _) = test_notebook_path(
             &path,
             Path::new("isort_expected.ipynb"),
             &settings::Settings::for_rule(Rule::UnsortedImports),
@@ -553,7 +553,7 @@ print("after empty cells")
     #[test]
     fn test_line_magics() -> Result<()> {
         let path = "line_magics.ipynb".to_string();
-        let (diagnostics, source_kind) = test_notebook_path(
+        let (diagnostics, source_kind, _) = test_notebook_path(
             &path,
             Path::new("line_magics_expected.ipynb"),
             &settings::Settings::for_rule(Rule::UnusedImport),
@@ -565,7 +565,7 @@ print("after empty cells")
     #[test]
     fn test_json_consistency() -> Result<()> {
         let path = "before_fix.ipynb".to_string();
-        let (_, source_kind) = test_notebook_path(
+        let (_, _, source_kind) = test_notebook_path(
             path,
             Path::new("after_fix.ipynb"),
             &settings::Settings::for_rule(Rule::UnusedImport),
diff --git a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
index a1e4c636f0ca6..e0ab3572a84d4 100644
--- a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
+++ b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__import_sorting.snap
@@ -47,18 +47,16 @@ isort.ipynb:cell 2:1:1: I001 [*] Import block is un-sorted or un-formatted
 7 9 | def foo():
 8 10 |     pass
 
-isort.ipynb:cell 2:6:1: I001 [*] Import block is un-sorted or un-formatted
-   |
- 4 |   def foo():
- 5 |       pass
- 6 | / from pathlib import Path
- 7 | | import sys
- 8 | | 
- 9 | | %matplotlib \
-   | |_^ I001
-10 |       --inline
-   |
-   = help: Organize imports
+isort.ipynb:cell 3:1:1: I001 [*] Import block is un-sorted or un-formatted
+  |
+1 | / from pathlib import Path
+2 | | import sys
+3 | | 
+4 | | %matplotlib \
+  | |_^ I001
+5 |       --inline
+  |
+  = help: Organize imports
 
 ℹ Fix
 6  6  | # Newline should be added here
@@ -71,12 +69,12 @@ isort.ipynb:cell 2:6:1: I001 [*] Import block is un-sorted or un-formatted
 12 12 | %matplotlib \
 13 13 |     --inline
 
-isort.ipynb:cell 3:5:1: I001 [*] Import block is un-sorted or un-formatted
+isort.ipynb:cell 3:7:1: I001 [*] Import block is un-sorted or un-formatted
   |
-3 |       --inline
-4 |   
-5 | / import math
-6 | | import abc
+5 |       --inline
+6 |   
+7 | / import math
+8 | | import abc
   |
   = help: Organize imports
 
diff --git a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
index 576b938b6b2e0..61211b2c18edb 100644
--- a/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
+++ b/crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap
@@ -7,6 +7,8 @@ line_magics.ipynb:cell 1:5:8: F401 [*] `os` imported but unused
 4 | 
 5 | import os
   |        ^^ F401
+6 | 
+7 | _ = math.pi
   |
   = help: Remove unused import: `os`
 
diff --git a/crates/ruff/src/test.rs b/crates/ruff/src/test.rs
index abf6c16131077..55ffec64260c1 100644
--- a/crates/ruff/src/test.rs
+++ b/crates/ruff/src/test.rs
@@ -62,8 +62,9 @@ pub(crate) fn test_notebook_path(
     path: impl AsRef<Path>,
     expected: impl AsRef<Path>,
     settings: &Settings,
-) -> Result<(Vec<Message>, SourceKind)> {
+) -> Result<(Vec<Message>, SourceKind, SourceKind)> {
     let mut source_kind = SourceKind::Jupyter(read_jupyter_notebook(path.as_ref())?);
+    let original_source_kind = source_kind.clone();
     let messages = test_contents(&mut source_kind, path.as_ref(), settings);
     let expected_notebook = read_jupyter_notebook(expected.as_ref())?;
     if let SourceKind::Jupyter(notebook) = &source_kind {
@@ -71,7 +72,7 @@ pub(crate) fn test_notebook_path(
         assert_eq!(notebook.index(), expected_notebook.index());
         assert_eq!(notebook.content(), expected_notebook.content());
     };
-    Ok((messages, source_kind))
+    Ok((messages, original_source_kind, source_kind))
 }
 
 /// Run [`check_path`] on a snippet of Python code.