diff --git a/bin/git-check-assertions b/bin/git-check-assertions
index 53fc277..fec0805 100755
--- a/bin/git-check-assertions
+++ b/bin/git-check-assertions
@@ -72,6 +72,9 @@ else
   echo "No main or master branch found." >&2
   exit 1
 fi
+restore() {
+  git restore -q --staged --worktree .
+}
 base="$(git merge-base "$base_branch" HEAD)"
 mapfile -t commits < <(git rev-list --reverse "${base}..HEAD")
 
@@ -92,12 +95,12 @@ for commit_hash in "${commits[@]}"; do
     printf '%s\n' "$block" | sed 's/^/> /'
     if ! bash -euo pipefail -c "$block"; then
       echo "git-check-assertions block failed in $commit_hash" >&2
-      git restore -q .
+      restore
       echo "Returning to $orig_ref"
       git checkout -q "$orig_ref"
       exit 1
     fi
-    git restore -q .
+    restore
   fi
   echo
 done
diff --git a/test/git-check-assertions.bats b/test/git-check-assertions.bats
index beb29ca..ee5097b 100755
--- a/test/git-check-assertions.bats
+++ b/test/git-check-assertions.bats
@@ -129,6 +129,37 @@ commit_with_assertion() {
   assert_size_zero readme
 }
 
+@test "should restore index when finished" {
+  git checkout -b feature
+  commit_with_assertion $'echo blah >> readme\ngit add readme'
+
+  run git-check-assertions
+
+  assert_success
+  assert_size_zero readme
+}
+
+@test "should restore index between commits" {
+  git checkout -b feature
+  commit_with_assertion $'echo blah >> readme\ngit add readme'
+  commit_with_assertion 'cp readme ../test'
+
+  run git-check-assertions
+
+  assert_success
+  assert_size_zero ../test
+}
+
+@test "should restore index after a failure" {
+  git checkout -b feature
+  commit_with_assertion $'echo blah >> readme\ngit add readme\nexit 3'
+
+  run git-check-assertions
+
+  assert_failure
+  assert_size_zero readme
+}
+
 @test "assertions should run against the version of the code inside of the commit" {
   git checkout -b feature
   echo commit1 >readme
@@ -157,6 +188,7 @@ commit_with_assertion() {
   run git-check-assertions
 
   assert_success
+  assert_file_contains src goodbye
 }
 
 @test "assert_success should succeed if the executed command succeeded" {