diff --git a/spec/migration_lock_timeout/migration_spec.rb b/spec/migration_lock_timeout/migration_spec.rb
index 080bc7c..aff25df 100644
--- a/spec/migration_lock_timeout/migration_spec.rb
+++ b/spec/migration_lock_timeout/migration_spec.rb
@@ -8,21 +8,12 @@
 ACTIVE_RECORD_MIGRATION_CLASS = ActiveRecord::Migration[ActiveRecord::VERSION::STRING.to_f]
 
 def expect_create_table
-  if ActiveRecord.gem_version >= '7.1'
-    expect(ActiveRecord::Base.connection).to receive(:execute).
-        with('BEGIN').
-        and_call_original
-    expect(ActiveRecord::Base.connection).to receive(:execute).
-      with(/CREATE TABLE/).
-      and_call_original
-  else
-    expect(ActiveRecord::Base.connection).to receive(:execute).
-        with('BEGIN', 'TRANSACTION').
-        and_call_original
-    expect(ActiveRecord::Base.connection).to receive(:execute).
-      with(/CREATE TABLE/).
-      and_call_original
-  end
+  expect(ActiveRecord::Base.connection).to receive(:execute).
+    with('BEGIN', 'TRANSACTION').
+    and_call_original
+  expect(ActiveRecord::Base.connection).to receive(:execute).
+    with(/CREATE TABLE/).
+    and_call_original
 end
 
 RSpec.describe ActiveRecord::Migration do
diff --git a/spec/migration_lock_timeout/spec_helper.rb b/spec/migration_lock_timeout/spec_helper.rb
index 51dc0f0..291edc8 100644
--- a/spec/migration_lock_timeout/spec_helper.rb
+++ b/spec/migration_lock_timeout/spec_helper.rb
@@ -18,4 +18,10 @@
       example.run
     end
   end
+
+  if ActiveRecord.gem_version >= '7.1'
+    config.after(:each) do |example|
+      ActiveRecord::Base.remove_connection
+    end
+  end
 end