From 7a79766a91f014ec78ccc10c28af491237eb1d58 Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Thu, 21 Jul 2022 23:32:42 +0200
Subject: [PATCH] Added improved doctest support

---
 CHANGELOG.md                                  |  1 +
 src/lib.rs                                    |  2 +-
 src/runtime.rs                                | 48 +++++++++++++++++--
 src/snapshots/doctest_runtime_rs__named.snap  | 10 ++++
 .../doctest_runtime_rs__unnamed_doctest.snap  |  8 ++++
 5 files changed, 63 insertions(+), 6 deletions(-)
 create mode 100644 src/snapshots/doctest_runtime_rs__named.snap
 create mode 100644 src/snapshots/doctest_runtime_rs__unnamed_doctest.snap

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 68ba7109..474ff362 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ All notable changes to insta and cargo-insta are documented here.
   quiet flag. This works around limitations with custom test harnesses
   such as cucumber.
 - Update RON to 0.7.1.
+- Added improved support for running insta from doctests. (#243)
 
 ## 1.15.0
 
diff --git a/src/lib.rs b/src/lib.rs
index 08b1d15c..8dd58afe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -36,7 +36,7 @@
 //!
 //! # Writing Tests
 //!
-//! ```no_run
+//! ```
 //! use insta::assert_debug_snapshot;
 //!
 //! #[test]
diff --git a/src/runtime.rs b/src/runtime.rs
index 95e0834d..7e34b726 100644
--- a/src/runtime.rs
+++ b/src/runtime.rs
@@ -75,11 +75,20 @@ pub enum ReferenceValue<'a> {
     Inline(&'a str),
 }
 
+fn is_doctest(function_name: &str) -> bool {
+    function_name.starts_with("rust_out::main::_doctest")
+}
+
 fn detect_snapshot_name(function_name: &str, module_path: &str) -> Result<String, &'static str> {
-    let name = Cow::Borrowed(function_name);
+    let mut name = function_name;
+
+    // simplify doctest names
+    if is_doctest(name) {
+        name = "unnamed_doctest";
+    }
 
     // clean test name first
-    let mut name = name.rsplit("::").next().unwrap();
+    name = name.rsplit("::").next().unwrap();
     let mut test_prefixed = false;
     if name.starts_with("test_") {
         name = &name[5..];
@@ -135,7 +144,9 @@ fn add_suffix_to_snapshot_name(name: Cow<'_, str>) -> Cow<'_, str> {
 }
 
 fn get_snapshot_filename(
+    function_name: &str,
     module_path: &str,
+    assertion_file: &str,
     snapshot_name: &str,
     cargo_workspace: &Path,
     base: &str,
@@ -149,7 +160,20 @@ fn get_snapshot_filename(
                 use std::fmt::Write;
                 let mut f = String::new();
                 if settings.prepend_module_to_snapshot() {
-                    write!(&mut f, "{}__", module_path.replace("::", "__")).unwrap();
+                    if is_doctest(function_name) {
+                        write!(
+                            &mut f,
+                            "doctest_{}__",
+                            Path::new(assertion_file)
+                                .file_name()
+                                .unwrap()
+                                .to_string_lossy()
+                                .replace('.', "_")
+                        )
+                        .unwrap();
+                    } else {
+                        write!(&mut f, "{}__", module_path.replace("::", "__")).unwrap();
+                    }
                 }
                 write!(
                     &mut f,
@@ -197,8 +221,14 @@ impl<'a> SnapshotAssertionContext<'a> {
                         .unwrap()
                         .into(),
                 };
-                let file =
-                    get_snapshot_filename(module_path, &name, &cargo_workspace, assertion_file);
+                let file = get_snapshot_filename(
+                    function_name,
+                    module_path,
+                    assertion_file,
+                    &name,
+                    &cargo_workspace,
+                    assertion_file,
+                );
                 if fs::metadata(&file).is_ok() {
                     old_snapshot = Some(Snapshot::from_file(&file)?);
                 }
@@ -456,3 +486,11 @@ pub fn assert_snapshot(
 
     Ok(())
 }
+
+/// Test snapshots in doctests.
+///
+/// ```
+/// insta::assert_yaml_snapshot!(vec![1, 2, 3]);
+/// insta::assert_yaml_snapshot!("named", vec![1, 2, 3, 4, 5]);
+/// ```
+const _DOCTEST: bool = false;
diff --git a/src/snapshots/doctest_runtime_rs__named.snap b/src/snapshots/doctest_runtime_rs__named.snap
new file mode 100644
index 00000000..de3eb3fc
--- /dev/null
+++ b/src/snapshots/doctest_runtime_rs__named.snap
@@ -0,0 +1,10 @@
+---
+source: src/runtime.rs
+expression: "vec![1, 2, 3, 4, 5]"
+---
+- 1
+- 2
+- 3
+- 4
+- 5
+
diff --git a/src/snapshots/doctest_runtime_rs__unnamed_doctest.snap b/src/snapshots/doctest_runtime_rs__unnamed_doctest.snap
new file mode 100644
index 00000000..bebca0a1
--- /dev/null
+++ b/src/snapshots/doctest_runtime_rs__unnamed_doctest.snap
@@ -0,0 +1,8 @@
+---
+source: src/runtime.rs
+expression: "vec![1, 2, 3]"
+---
+- 1
+- 2
+- 3
+