diff --git a/tests/replace_tests.yaml b/tests/replace_tests.yaml
new file mode 100644
index 0000000..9cb5fde
--- /dev/null
+++ b/tests/replace_tests.yaml
@@ -0,0 +1,42 @@
+- original: |
+    def one():
+        pass
+
+    def two():
+        "Two"
+  stdin: |
+    def two(): pass
+  args: ["two", "--replace"]
+  expected: |
+    def one():
+        pass
+
+    def two(): pass
+
+- original: |
+    import os
+
+    @decorated
+    def one():
+        "This has multiple lines and a decorator"
+        return 1 + 2
+
+    def two():
+        "Two"
+  stdin: |
+    def one():
+        # No decorator now, but multiple lines
+        a = 1 + 2
+        return a * 3
+  args: ["one", "--replace"]
+  expected: |
+    import os
+
+    def one():
+        # No decorator now, but multiple lines
+        a = 1 + 2
+        return a * 3
+
+
+    def two():
+        "Two"
diff --git a/tests/test_replace.py b/tests/test_replace.py
index 3c328f8..7d414b2 100644
--- a/tests/test_replace.py
+++ b/tests/test_replace.py
@@ -2,8 +2,28 @@
 from symbex.cli import cli
 import pathlib
 import pytest
-import sys
 from unittest.mock import patch
+import yaml
+
+
+@pytest.mark.parametrize(
+    "test", yaml.safe_load(open(pathlib.Path(__file__).parent / "replace_tests.yaml"))
+)
+def test_replace(test):
+    original, stdin, args, expected = (
+        test["original"],
+        test["stdin"],
+        test["args"],
+        test["expected"],
+    )
+    runner = CliRunner()
+    with runner.isolated_filesystem() as root:
+        path = pathlib.Path(root) / "code.py"
+        path.write_text(original, "utf-8")
+        result = runner.invoke(cli, args, input=stdin, catch_exceptions=False)
+        modified = path.read_text("utf-8")
+    assert result.exit_code == 0
+    assert modified.strip() == expected.strip()
 
 
 @pytest.mark.parametrize(
@@ -39,7 +59,6 @@ def test_replace_errors(files, args, error):
             for path, code in files.items():
                 (root / path).parent.mkdir(parents=True, exist_ok=True)
                 (root / path).write_text(code, "utf-8")
-            runner = CliRunner()
             result = runner.invoke(cli, args + ["--replace"], catch_exceptions=False)
     assert result.exit_code == 1
     assert result.output.strip() == error