diff --git a/test/test.rb b/test/test.rb
index 872820c6b13d336a4101033978b0b690d680fa8b..6cda7a395152df157b1ab10483d64085f863885c 100644
--- a/test/test.rb
+++ b/test/test.rb
@@ -219,18 +219,14 @@ module BrakemanTester::RescanTestHelper
     replace_with_sexp file do |parsed|
       class_body = parsed.body
 
-      if class_body[1].node_type == :block
-        class_body[1].reject! do |node|
-          node.is_a? Sexp and
-          node.node_type == :defn and
-          node.method_name == method_name
-        end
-      elsif class_body[1].node_type == :defn and
-        class_body[1].method_name == method_name
-
-        class_body.delete_at 1
+      class_body.reject! do |node|
+        node.is_a? Sexp and
+        node.node_type == :defn and
+        node.method_name == method_name
       end
 
+      parsed.body = class_body
+
       parsed
     end
   end