diff --git a/codemod/src/main.rs b/codemod/src/main.rs
index a1ab50f..d8cbc8b 100644
--- a/codemod/src/main.rs
+++ b/codemod/src/main.rs
@@ -1,3 +1,5 @@
+mod tests;
+
 use regex::Regex;
 use rnix::ast::{AstToken, AttrpathValue, Comment, Expr, Lambda, Param};
 use rnix::{SyntaxKind, SyntaxNode, SyntaxToken};
@@ -169,7 +171,7 @@ fn parse_doc_comment(
     for line in raw.lines() {
         let mut line = line.trim_end();
 
-        let trimmed = line.clone().trim();
+        let trimmed = line.trim();
 
         if trimmed.starts_with("Type:") {
             state = ParseState::Type;
@@ -274,7 +276,7 @@ fn get_binding_name(token: &SyntaxToken) -> Option<String> {
     name
 }
 
-fn get_argument_docs(token: &SyntaxToken, ident: &str) -> Option<String> {
+fn get_argument_docs(token: &SyntaxToken, _ident: &str) -> Option<String> {
     let mut step = token.next_sibling_or_token();
 
     // Find the Expr that is a lambda.
@@ -464,7 +466,7 @@ fn main() {
             let f_name = entry.file_name().to_string_lossy();
 
             if f_name.ends_with(".nix") {
-                modify_file(entry.path().to_path_buf());
+                modify_file_inplace(entry.path().to_path_buf());
             }
         }
     } else {
@@ -472,7 +474,28 @@ fn main() {
     }
 }
 
-fn modify_file(file_path: PathBuf) {
+fn replace_all(syntax: SyntaxNode) -> (Option<SyntaxNode>, i32) {
+    let mut maybe_replaced = replace_first_comment(&syntax);
+    let mut count = 0;
+    let r: Option<SyntaxNode> = loop {
+        if let Some(replaced) = maybe_replaced {
+            // Maybe we can replace more
+            count += 1;
+            let result = replace_first_comment(&replaced);
+
+            // If we cannot replace more comments
+            if result.is_none() {
+                break Some(replaced);
+            }
+            maybe_replaced = result;
+        } else {
+            break None;
+        }
+    };
+    (r, count)
+}
+
+fn modify_file_inplace(file_path: PathBuf) -> () {
     let contents = fs::read_to_string(&file_path);
     if let Err(_) = contents {
         println!("Could not read the file {:?}", file_path);
@@ -492,28 +515,9 @@ fn modify_file(file_path: PathBuf) {
         return;
     }
 
-    // println!("Syntax \n {:?}", root.clone().unwrap().syntax().clone());
     let syntax = root.unwrap().syntax().clone_for_update();
-    let mut maybe_replaced = replace_first_comment(&syntax);
-    let mut count = 0;
-    let r: Option<SyntaxNode> = loop {
-        if let Some(replaced) = maybe_replaced {
-            // Maybe we can replace more
-            count += 1;
-            let result = replace_first_comment(&replaced);
-
-            // If we cannot replace more comments
-            if result.is_none() {
-                break Some(replaced);
-            }
-            maybe_replaced = result;
-        } else {
-            break None;
-        }
-    };
-
     let display_name = file_path.to_str().unwrap();
-    if let Some(updates) = r {
+    if let (Some(updates), count) = replace_all(syntax) {
         let mut file = File::create(&file_path).unwrap();
         file.write_all(updates.text().to_string().as_bytes()).ok();
         println!("{display_name} - Changed {count} comments");
diff --git a/codemod/src/tests.rs b/codemod/src/tests.rs
new file mode 100644
index 0000000..f82ff78
--- /dev/null
+++ b/codemod/src/tests.rs
@@ -0,0 +1,23 @@
+#[cfg(test)]
+mod tests {
+    use std::{fs, path::PathBuf};
+
+    use crate::replace_all;
+    use expect_test::expect_file;
+    use rowan::ast::AstNode;
+
+    #[test]
+    fn test_snapshot_mixed() {
+        let path: PathBuf = [env!("CARGO_MANIFEST_DIR"), "test_data", "mixed.nix"]
+            .iter()
+            .collect();
+        let src = fs::read_to_string(&path).unwrap();
+        let nix = rnix::Root::parse(&src).ok().expect("failed to parse input");
+
+        let (res, _count) = replace_all(nix.syntax().to_owned());
+
+        let actual = res.unwrap().to_string();
+
+        expect_file![path.with_extension("expect")].assert_eq(&actual);
+    }
+}
diff --git a/codemod/test/args.nix b/codemod/test_data/args.nix
similarity index 100%
rename from codemod/test/args.nix
rename to codemod/test_data/args.nix
diff --git a/codemod/test/args/default.nix b/codemod/test_data/args/default.nix
similarity index 100%
rename from codemod/test/args/default.nix
rename to codemod/test_data/args/default.nix
diff --git a/codemod/test_data/mixed.expect b/codemod/test_data/mixed.expect
new file mode 100644
index 0000000..072bd64
--- /dev/null
+++ b/codemod/test_data/mixed.expect
@@ -0,0 +1,64 @@
+{
+  /**
+    Short tagline
+    in multiline
+    with empty lines
+
+    indepentent block of lines
+    should not be indented differently
+
+    > Note: this is a note.
+
+    ```nix
+    {
+      a = {
+        b = {
+          c = {
+            d = 1;
+          };
+        };
+      };
+    }
+    ```
+
+    Codeblocks must keep indentation.
+
+
+    # Inputs
+
+    `dir`
+
+    : 1\. Function argument
+
+    # Type
+
+    ```
+    Indent :: a ->
+    {
+      a: {
+        b: {
+          c: {
+            d: Int;
+          };
+        };
+      };
+    }
+    ```
+
+    # Examples
+    :::{.example}
+    ## `packageDir` usage example
+
+    ```nix
+    packageDir "foo"
+    => {
+      foo = {
+        src = ./.;
+      };
+    }
+    ```
+
+    :::
+  */
+  packageDir = dir: 1;
+}
diff --git a/codemod/test/mixed.nix b/codemod/test_data/mixed.nix
similarity index 84%
rename from codemod/test/mixed.nix
rename to codemod/test_data/mixed.nix
index 973722e..55ccde6 100644
--- a/codemod/test/mixed.nix
+++ b/codemod/test_data/mixed.nix
@@ -3,6 +3,11 @@
     in multiline
     with empty lines
 
+    indepentent block of lines
+    should not be indented differently
+
+    > Note: this is a note.
+
     ```nix
     {
       a = {
diff --git a/codemod/test/nested/sample.nix b/codemod/test_data/nested/sample.nix
similarity index 100%
rename from codemod/test/nested/sample.nix
rename to codemod/test_data/nested/sample.nix
diff --git a/codemod/test/simple.nix b/codemod/test_data/simple.nix
similarity index 100%
rename from codemod/test/simple.nix
rename to codemod/test_data/simple.nix