From db700a93d49aad3b47b6151cd140f6cd43a50b63 Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:52:31 -0400
Subject: [PATCH 1/8] New helper module, check for within migration

---
 lib/rubocop-sequel.rb                         |  2 +
 lib/rubocop/cop/sequel/helpers/migration.rb   | 25 +++++++++
 .../cop/sequel/irreversible_migration.rb      |  7 ++-
 .../cop/sequel/irreversible_migration_spec.rb | 55 +++++++++++++------
 4 files changed, 70 insertions(+), 19 deletions(-)
 create mode 100644 lib/rubocop/cop/sequel/helpers/migration.rb

diff --git a/lib/rubocop-sequel.rb b/lib/rubocop-sequel.rb
index 1c2ed2f..1113391 100644
--- a/lib/rubocop-sequel.rb
+++ b/lib/rubocop-sequel.rb
@@ -7,6 +7,8 @@
 
 RuboCop::Sequel::Inject.defaults!
 
+require 'rubocop/cop/sequel/helpers/migration'
+
 require 'rubocop/cop/sequel/concurrent_index'
 require 'rubocop/cop/sequel/irreversible_migration'
 require 'rubocop/cop/sequel/json_column'
diff --git a/lib/rubocop/cop/sequel/helpers/migration.rb b/lib/rubocop/cop/sequel/helpers/migration.rb
new file mode 100644
index 0000000..68adf15
--- /dev/null
+++ b/lib/rubocop/cop/sequel/helpers/migration.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module RuboCop
+  module Cop
+    module Sequel
+      module Helpers
+        # Migration contains helper methods for detecting if a node is inside a `Sequel.migration` block
+        module Migration
+          extend NodePattern::Macros
+
+          def_node_matcher :sequel_migration_block?, <<~MATCHER
+            (block
+              (send
+                (const nil? :Sequel) :migration ...)
+              ...)
+          MATCHER
+
+          def within_sequel_migration?(node)
+            node.each_ancestor(:block).any? { |ancestor| sequel_migration_block?(ancestor) }
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/rubocop/cop/sequel/irreversible_migration.rb b/lib/rubocop/cop/sequel/irreversible_migration.rb
index d3455c1..3f794fd 100644
--- a/lib/rubocop/cop/sequel/irreversible_migration.rb
+++ b/lib/rubocop/cop/sequel/irreversible_migration.rb
@@ -5,6 +5,8 @@ module Cop
     module Sequel
       # IrreversibleMigration looks for methods inside a `change` block that cannot be reversed.
       class IrreversibleMigration < Base
+        include RuboCop::Cop::Sequel::Helpers::Migration
+
         # https://sequel.jeremyevans.net/rdoc/files/doc/migration_rdoc.html#label-A+Basic+Migration
         VALID_CHANGE_METHODS = %i[
           create_table
@@ -26,11 +28,12 @@ class IrreversibleMigration < Base
           set_column_allow_null
         ].freeze
 
-        MSG = 'Avoid using `%<name>s` inside a `change` block. Use `up` & `down` blocks instead.'
-        PRIMARY_KEY_MSG = 'Avoid using `add_primary_key` with an array argument inside a `change` block.'
+        MSG = 'Avoid using "%<name>s" inside a "change" block. Use "up" & "down" blocks instead.'
+        PRIMARY_KEY_MSG = 'Avoid using "add_primary_key" with an array argument inside a "change" block.'
 
         def on_block(node)
           return unless node.method_name == :change
+          return unless within_sequel_migration?(node)
 
           body = node.body
           return unless body
diff --git a/spec/rubocop/cop/sequel/irreversible_migration_spec.rb b/spec/rubocop/cop/sequel/irreversible_migration_spec.rb
index 3bbb57c..ac4a64b 100644
--- a/spec/rubocop/cop/sequel/irreversible_migration_spec.rb
+++ b/spec/rubocop/cop/sequel/irreversible_migration_spec.rb
@@ -6,10 +6,12 @@
   context 'when inside a change block' do
     let(:invalid_source) do
       <<~SOURCE
-        change do
-          alter_table(:stores) do
-            drop_column(:products, :name)
-            drop_index(:products, :price)
+        Sequel.migration do
+          change do
+            alter_table(:stores) do
+              drop_column(:products, :name)
+              drop_index(:products, :price)
+            end
           end
         end
       SOURCE
@@ -17,11 +19,13 @@
 
     let(:valid_source) do
       <<~SOURCE
-        change do
-          alter_table(:stores) do
-            add_primary_key(:id)
-            add_column(:products, :name)
-            add_index(:products, :price)
+        Sequel.migration do
+          change do
+            alter_table(:stores) do
+              add_primary_key(:id)
+              add_column(:products, :name)
+              add_index(:products, :price)
+            end
           end
         end
       SOURCE
@@ -40,9 +44,11 @@
     describe 'and an array is passed into `add_primary_key`' do
       let(:source) do
         <<~SOURCE
-          change do
-            alter_table(:stores) do
-              add_primary_key([:owner_id, :name])
+          Sequel.migration do
+            change do
+              alter_table(:stores) do
+                add_primary_key([:owner_id, :name])
+              end
             end
           end
         SOURCE
@@ -58,11 +64,13 @@
   context 'when inside an up block' do
     let(:source) do
       <<~SOURCE
-        up do
-          alter_table(:stores) do
-            add_primary_key([:owner_id, :name])
-            add_column(:products, :name)
-            drop_index(:products, :price)
+        Sequel.migration do
+          up do
+            alter_table(:stores) do
+              add_primary_key([:owner_id, :name])
+              add_column(:products, :name)
+              drop_index(:products, :price)
+            end
           end
         end
       SOURCE
@@ -73,4 +81,17 @@
       expect(offenses).to be_empty
     end
   end
+
+  context 'when a change block is used outside of a Sequel migration' do
+    let(:source) do
+      <<~SOURCE
+        it { expect { subject }.to change { document_count(user_id) }.by(-1) }
+      SOURCE
+    end
+
+    it 'does not register an offense with any methods' do
+      offenses = inspect_source(source)
+      expect(offenses).to be_empty
+    end
+  end
 end

From 0b87e4e45d384b02a7ba0695a7bb4c2f3a431760 Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:57:42 -0400
Subject: [PATCH 2/8] Version number + config update

---
 config/default.yml            | 10 ++++++++--
 lib/rubocop/sequel/version.rb |  2 +-
 rubocop-sequel.gemspec        |  2 +-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/config/default.yml b/config/default.yml
index d2c2396..05c328c 100644
--- a/config/default.yml
+++ b/config/default.yml
@@ -16,8 +16,8 @@ Sequel/IrreversibleMigration:
     Warns against using certain methods inside a migration file's `change` block that would cause the migration to become irreversible.
   Reference: https://www.rubydoc.info/gems/rubocop-sequel/RuboCop/Cop/Sequel/IrreversibleMigration
   Enabled: true
-  VersionAdded: 0.3.4
-  VersionChanged: 0.3.4
+  VersionAdded: 0.3.5
+  VersionChanged: 0.3.6
 
 Sequel/JSONColumn:
   Description: >-
@@ -52,3 +52,9 @@ Sequel/SaveChanges:
   SafeAutoCorrect: false
   VersionAdded: '0.0.1'
   VersionChanged: '0.3.3'
+
+Sequel/Helpers/Migration:
+  Description: Contains helper methods for detecting if a node is inside a `Sequel.migration` block
+  Reference: https://www.rubydoc.info/gems/rubocop-sequel/RuboCop/Cop/Sequel/Helpers/Migration
+  VersionAdded: 0.3.6
+  VersionChanged: 0.3.6
diff --git a/lib/rubocop/sequel/version.rb b/lib/rubocop/sequel/version.rb
index 4dff340..8e30804 100644
--- a/lib/rubocop/sequel/version.rb
+++ b/lib/rubocop/sequel/version.rb
@@ -4,7 +4,7 @@ module RuboCop
   module Sequel
     # This module holds the RuboCop Sequel version information.
     module Version
-      STRING = '0.3.5'
+      STRING = '0.3.6'
     end
   end
 end
diff --git a/rubocop-sequel.gemspec b/rubocop-sequel.gemspec
index 35b5a4c..c41ae8a 100644
--- a/rubocop-sequel.gemspec
+++ b/rubocop-sequel.gemspec
@@ -12,7 +12,7 @@ Gem::Specification.new do |gem|
   gem.executables   = gem.files.grep(%r{^bin/}).map { |f| File.basename(f) }
   gem.name          = 'rubocop-sequel'
   gem.require_paths = ['lib']
-  gem.version       = '0.3.5'
+  gem.version       = '0.3.6'
   gem.metadata['rubygems_mfa_required'] = 'true'
 
   gem.required_ruby_version = '>= 2.5'

From b96d6505e2cef9160598e781d354fcddfb5f217e Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Thu, 24 Oct 2024 13:11:39 -0400
Subject: [PATCH 3/8] Project spec fix (alphabetical order)

---
 config/default.yml | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/config/default.yml b/config/default.yml
index 05c328c..1e25fcb 100644
--- a/config/default.yml
+++ b/config/default.yml
@@ -11,6 +11,12 @@ Sequel/ConcurrentIndex:
   VersionAdded: 0.0.1
   VersionChanged: 0.3.3
 
+Sequel/Helpers/Migration:
+  Description: Contains helper methods for detecting if a node is inside a `Sequel.migration` block
+  Reference: https://www.rubydoc.info/gems/rubocop-sequel/RuboCop/Cop/Sequel/Helpers/Migration
+  VersionAdded: 0.3.6
+  VersionChanged: 0.3.6
+
 Sequel/IrreversibleMigration:
   Description: >-
     Warns against using certain methods inside a migration file's `change` block that would cause the migration to become irreversible.
@@ -52,9 +58,3 @@ Sequel/SaveChanges:
   SafeAutoCorrect: false
   VersionAdded: '0.0.1'
   VersionChanged: '0.3.3'
-
-Sequel/Helpers/Migration:
-  Description: Contains helper methods for detecting if a node is inside a `Sequel.migration` block
-  Reference: https://www.rubydoc.info/gems/rubocop-sequel/RuboCop/Cop/Sequel/Helpers/Migration
-  VersionAdded: 0.3.6
-  VersionChanged: 0.3.6

From d48cbaf21d200aa1d6dbdd63706f8bb4ecd8e857 Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Fri, 25 Oct 2024 10:26:24 -0400
Subject: [PATCH 4/8] Fix uninitialized constant issue

---
 lib/rubocop-sequel.rb                            | 2 +-
 lib/rubocop/cop/sequel/irreversible_migration.rb | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/rubocop-sequel.rb b/lib/rubocop-sequel.rb
index 1113391..d2e2929 100644
--- a/lib/rubocop-sequel.rb
+++ b/lib/rubocop-sequel.rb
@@ -7,7 +7,7 @@
 
 RuboCop::Sequel::Inject.defaults!
 
-require 'rubocop/cop/sequel/helpers/migration'
+require_relative 'rubocop/cop/sequel/helpers/migration'
 
 require 'rubocop/cop/sequel/concurrent_index'
 require 'rubocop/cop/sequel/irreversible_migration'
diff --git a/lib/rubocop/cop/sequel/irreversible_migration.rb b/lib/rubocop/cop/sequel/irreversible_migration.rb
index 3f794fd..3443fd8 100644
--- a/lib/rubocop/cop/sequel/irreversible_migration.rb
+++ b/lib/rubocop/cop/sequel/irreversible_migration.rb
@@ -5,7 +5,7 @@ module Cop
     module Sequel
       # IrreversibleMigration looks for methods inside a `change` block that cannot be reversed.
       class IrreversibleMigration < Base
-        include RuboCop::Cop::Sequel::Helpers::Migration
+        include Helpers::Migration
 
         # https://sequel.jeremyevans.net/rdoc/files/doc/migration_rdoc.html#label-A+Basic+Migration
         VALID_CHANGE_METHODS = %i[

From b445c3c4563c8a85537a5bd556f593980abe8b11 Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Fri, 25 Oct 2024 10:26:32 -0400
Subject: [PATCH 5/8] Remove unnecessary entry

---
 config/default.yml | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/config/default.yml b/config/default.yml
index 1e25fcb..6fdc2d9 100644
--- a/config/default.yml
+++ b/config/default.yml
@@ -11,12 +11,6 @@ Sequel/ConcurrentIndex:
   VersionAdded: 0.0.1
   VersionChanged: 0.3.3
 
-Sequel/Helpers/Migration:
-  Description: Contains helper methods for detecting if a node is inside a `Sequel.migration` block
-  Reference: https://www.rubydoc.info/gems/rubocop-sequel/RuboCop/Cop/Sequel/Helpers/Migration
-  VersionAdded: 0.3.6
-  VersionChanged: 0.3.6
-
 Sequel/IrreversibleMigration:
   Description: >-
     Warns against using certain methods inside a migration file's `change` block that would cause the migration to become irreversible.

From 880aac6ea185f5d61c9dfc5675d259fb8fe8bf72 Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Fri, 25 Oct 2024 10:47:49 -0400
Subject: [PATCH 6/8] New spec helper: Migration

---
 spec/rubocop/cop/sequel/helpers/migration.rb  | 19 +++++++
 .../cop/sequel/irreversible_migration_spec.rb | 52 ++++++++-----------
 spec/spec_helper.rb                           |  2 +
 3 files changed, 44 insertions(+), 29 deletions(-)
 create mode 100644 spec/rubocop/cop/sequel/helpers/migration.rb

diff --git a/spec/rubocop/cop/sequel/helpers/migration.rb b/spec/rubocop/cop/sequel/helpers/migration.rb
new file mode 100644
index 0000000..3f1cd66
--- /dev/null
+++ b/spec/rubocop/cop/sequel/helpers/migration.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Spec
+  module Helpers
+    module Migration
+      extend CopHelper
+
+      def inspect_source_within_migration(source, file = nil)
+        migration_wrapped_source = <<~SOURCE
+          Sequel.migration do
+            #{source}
+          end
+        SOURCE
+
+        inspect_source(migration_wrapped_source, file)
+      end
+    end
+  end
+end
diff --git a/spec/rubocop/cop/sequel/irreversible_migration_spec.rb b/spec/rubocop/cop/sequel/irreversible_migration_spec.rb
index ac4a64b..144c384 100644
--- a/spec/rubocop/cop/sequel/irreversible_migration_spec.rb
+++ b/spec/rubocop/cop/sequel/irreversible_migration_spec.rb
@@ -1,17 +1,17 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::Cop::Sequel::IrreversibleMigration do
+  include Spec::Helpers::Migration
+
   subject(:cop) { described_class.new }
 
   context 'when inside a change block' do
     let(:invalid_source) do
       <<~SOURCE
-        Sequel.migration do
-          change do
-            alter_table(:stores) do
-              drop_column(:products, :name)
-              drop_index(:products, :price)
-            end
+        change do
+          alter_table(:stores) do
+            drop_column(:products, :name)
+            drop_index(:products, :price)
           end
         end
       SOURCE
@@ -19,43 +19,39 @@
 
     let(:valid_source) do
       <<~SOURCE
-        Sequel.migration do
-          change do
-            alter_table(:stores) do
-              add_primary_key(:id)
-              add_column(:products, :name)
-              add_index(:products, :price)
-            end
+        change do
+          alter_table(:stores) do
+            add_primary_key(:id)
+            add_column(:products, :name)
+            add_index(:products, :price)
           end
         end
       SOURCE
     end
 
     it 'registers an offense when there is an invalid method' do
-      offenses = inspect_source(invalid_source)
+      offenses = inspect_source_within_migration(invalid_source)
       expect(offenses.size).to eq(2)
     end
 
     it 'does not register an offense with valid methods' do
-      offenses = inspect_source(valid_source)
+      offenses = inspect_source_within_migration(valid_source)
       expect(offenses).to be_empty
     end
 
     describe 'and an array is passed into `add_primary_key`' do
       let(:source) do
         <<~SOURCE
-          Sequel.migration do
-            change do
-              alter_table(:stores) do
-                add_primary_key([:owner_id, :name])
-              end
+          change do
+            alter_table(:stores) do
+              add_primary_key([:owner_id, :name])
             end
           end
         SOURCE
       end
 
       it 'registers an offense' do
-        offenses = inspect_source(source)
+        offenses = inspect_source_within_migration(source)
         expect(offenses.size).to eq(1)
       end
     end
@@ -64,20 +60,18 @@
   context 'when inside an up block' do
     let(:source) do
       <<~SOURCE
-        Sequel.migration do
-          up do
-            alter_table(:stores) do
-              add_primary_key([:owner_id, :name])
-              add_column(:products, :name)
-              drop_index(:products, :price)
-            end
+        up do
+          alter_table(:stores) do
+            add_primary_key([:owner_id, :name])
+            add_column(:products, :name)
+            drop_index(:products, :price)
           end
         end
       SOURCE
     end
 
     it 'does not register an offense with any methods' do
-      offenses = inspect_source(source)
+      offenses = inspect_source_within_migration(source)
       expect(offenses).to be_empty
     end
   end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index a206132..b5ad7e5 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -4,6 +4,8 @@
 require 'rubocop/rspec/support'
 require 'rubocop-sequel'
 
+require_relative 'rubocop/cop/sequel/helpers/migration'
+
 RSpec.configure do |config|
   config.include RuboCop::RSpec::ExpectOffense
   config.disable_monkey_patching!

From 07485ab9baa268e8ca0aaf7aaa2a334c39a32671 Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Fri, 25 Oct 2024 10:47:57 -0400
Subject: [PATCH 7/8] Add .DS_Store to gitignore

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index b831235..07a403b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 .tags
 *.gem
 Gemfile.lock
+.DS_Store

From e6a154ee478425b7d46fd8b35fd890291c8fbbe3 Mon Sep 17 00:00:00 2001
From: Alasdair Wallace Mackie
 <73504675+AlasdairWallaceMackie@users.noreply.github.com>
Date: Fri, 25 Oct 2024 12:01:08 -0400
Subject: [PATCH 8/8] Cops will check for migration block, refactor specs

---
 lib/rubocop/cop/sequel/concurrent_index.rb     |  4 ++++
 lib/rubocop/cop/sequel/json_column.rb          |  4 ++++
 lib/rubocop/cop/sequel/partial_constraint.rb   |  3 +++
 .../cop/sequel/concurrent_index_spec.rb        | 10 ++++++----
 spec/rubocop/cop/sequel/json_column_spec.rb    | 18 ++++++++++--------
 .../cop/sequel/partial_constraint_spec.rb      |  4 +++-
 6 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/lib/rubocop/cop/sequel/concurrent_index.rb b/lib/rubocop/cop/sequel/concurrent_index.rb
index b16e92f..bd22cb0 100644
--- a/lib/rubocop/cop/sequel/concurrent_index.rb
+++ b/lib/rubocop/cop/sequel/concurrent_index.rb
@@ -5,6 +5,8 @@ module Cop
     module Sequel
       # ConcurrentIndex looks for non-concurrent index creation.
       class ConcurrentIndex < Base
+        include Helpers::Migration
+
         MSG = 'Specify `concurrently` option when creating or dropping an index.'
         RESTRICT_ON_SEND = %i[add_index drop_index].freeze
 
@@ -13,6 +15,8 @@ class ConcurrentIndex < Base
         MATCHER
 
         def on_send(node)
+          return unless within_sequel_migration?(node)
+
           indexes?(node) do |args|
             add_offense(node.loc.selector, message: MSG) if offensive?(args)
           end
diff --git a/lib/rubocop/cop/sequel/json_column.rb b/lib/rubocop/cop/sequel/json_column.rb
index 2e98b38..e049df0 100644
--- a/lib/rubocop/cop/sequel/json_column.rb
+++ b/lib/rubocop/cop/sequel/json_column.rb
@@ -5,6 +5,8 @@ module Cop
     module Sequel
       # JSONColumn looks for non-JSONB columns.
       class JSONColumn < Base
+        include Helpers::Migration
+
         MSG = 'Use JSONB rather than JSON or hstore'
         RESTRICT_ON_SEND = %i[add_column].freeze
 
@@ -22,12 +24,14 @@ class JSONColumn < Base
 
         def on_send(node)
           return unless json_or_hstore?(node)
+          return unless within_sequel_migration?(node)
 
           add_offense(node.loc.selector, message: MSG)
         end
 
         def on_block(node)
           return unless node.send_node.method_name == :create_table
+          return unless within_sequel_migration?(node)
 
           node.each_node(:send) do |method|
             next unless column_method?(method) || column_type?(method)
diff --git a/lib/rubocop/cop/sequel/partial_constraint.rb b/lib/rubocop/cop/sequel/partial_constraint.rb
index baf3ab4..b4d176b 100644
--- a/lib/rubocop/cop/sequel/partial_constraint.rb
+++ b/lib/rubocop/cop/sequel/partial_constraint.rb
@@ -5,6 +5,8 @@ module Cop
     module Sequel
       # PartialConstraint looks for missed usage of partial indexes.
       class PartialConstraint < Base
+        include Helpers::Migration
+
         MSG = "Constraint can't be partial, use where argument with index"
         RESTRICT_ON_SEND = %i[add_unique_constraint].freeze
 
@@ -14,6 +16,7 @@ class PartialConstraint < Base
 
         def on_send(node)
           return unless add_partial_constraint?(node)
+          return unless within_sequel_migration?(node)
 
           add_offense(node.loc.selector, message: MSG)
         end
diff --git a/spec/rubocop/cop/sequel/concurrent_index_spec.rb b/spec/rubocop/cop/sequel/concurrent_index_spec.rb
index d1b64f9..0978acd 100644
--- a/spec/rubocop/cop/sequel/concurrent_index_spec.rb
+++ b/spec/rubocop/cop/sequel/concurrent_index_spec.rb
@@ -1,11 +1,13 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::Cop::Sequel::ConcurrentIndex do
+  include Spec::Helpers::Migration
+
   subject(:cop) { described_class.new }
 
   context 'without the concurrent option' do
     it 'registers an offense without options' do
-      offenses = inspect_source(<<~SOURCE)
+      offenses = inspect_source_within_migration(<<~SOURCE)
         add_index(:products, :name)
         drop_index(:products, :name)
       SOURCE
@@ -13,7 +15,7 @@
     end
 
     it 'registers an offense with other options' do
-      offenses = inspect_source(<<~SOURCE)
+      offenses = inspect_source_within_migration(<<~SOURCE)
         add_index(:products, :name, unique: true)
         drop_index(:products, :name, unique: true)
       SOURCE
@@ -21,7 +23,7 @@
     end
 
     it 'registers an offense with composite index' do
-      offenses = inspect_source(<<~SOURCE)
+      offenses = inspect_source_within_migration(<<~SOURCE)
         add_index(:products, [:name, :price], unique: true)
         drop_index(:products, [:name, :price])
       SOURCE
@@ -30,7 +32,7 @@
   end
 
   it 'does not register an offense when using concurrent option' do
-    offenses = inspect_source(<<~SOURCE)
+    offenses = inspect_source_within_migration(<<~SOURCE)
       add_index(:products, :name, unique: true, concurrently: true)
       drop_index(:products, :name, concurrently: true)
     SOURCE
diff --git a/spec/rubocop/cop/sequel/json_column_spec.rb b/spec/rubocop/cop/sequel/json_column_spec.rb
index e9c3d73..6014137 100644
--- a/spec/rubocop/cop/sequel/json_column_spec.rb
+++ b/spec/rubocop/cop/sequel/json_column_spec.rb
@@ -1,48 +1,50 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::Cop::Sequel::JSONColumn do
+  include Spec::Helpers::Migration
+
   subject(:cop) { described_class.new }
 
   context 'with add_column' do
     it 'registers an offense when using json type' do
-      offenses = inspect_source('add_column(:products, :type, :json)')
+      offenses = inspect_source_within_migration('add_column(:products, :type, :json)')
       expect(offenses.size).to eq(1)
     end
 
     it 'registers an offense when using hstore type' do
-      offenses = inspect_source('add_column(:products, :type, :hstore)')
+      offenses = inspect_source_within_migration('add_column(:products, :type, :hstore)')
       expect(offenses.size).to eq(1)
     end
 
     it 'does not register an offense when using jsonb' do
-      offenses = inspect_source('add_column(:products, :type, :jsonb)')
+      offenses = inspect_source_within_migration('add_column(:products, :type, :jsonb)')
       expect(offenses).to be_empty
     end
   end
 
   context 'with create_table' do
     it 'registers an offense when using json as a method' do
-      offenses = inspect_source('create_table(:products) { json :type, default: {} }')
+      offenses = inspect_source_within_migration('create_table(:products) { json :type, default: {} }')
       expect(offenses.size).to eq(1)
     end
 
     it 'registers an offense when using the column method with hstore' do
-      offenses = inspect_source('create_table(:products) { column :type, :hstore }')
+      offenses = inspect_source_within_migration('create_table(:products) { column :type, :hstore }')
       expect(offenses.size).to eq(1)
     end
 
     it 'does not register an offense when using jsonb as column type`' do
-      offenses = inspect_source('create_table(:products) { column :type, :jsonb }')
+      offenses = inspect_source_within_migration('create_table(:products) { column :type, :jsonb }')
       expect(offenses).to be_empty
     end
 
     it 'does not register an offense when using jsonb' do
-      offenses = inspect_source('create_table(:products) { jsonb :type }')
+      offenses = inspect_source_within_migration('create_table(:products) { jsonb :type }')
       expect(offenses).to be_empty
     end
 
     it 'does not register an offense when using a simple type' do
-      offenses = inspect_source('create_table(:products) { integer :type, default: 0 }')
+      offenses = inspect_source_within_migration('create_table(:products) { integer :type, default: 0 }')
       expect(offenses).to be_empty
     end
   end
diff --git a/spec/rubocop/cop/sequel/partial_constraint_spec.rb b/spec/rubocop/cop/sequel/partial_constraint_spec.rb
index e19391c..128e885 100644
--- a/spec/rubocop/cop/sequel/partial_constraint_spec.rb
+++ b/spec/rubocop/cop/sequel/partial_constraint_spec.rb
@@ -1,10 +1,12 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::Cop::Sequel::PartialConstraint do
+  include Spec::Helpers::Migration
+
   subject(:cop) { described_class.new }
 
   it 'registers an offense when using where for constraint' do
-    offenses = inspect_source(<<~RUBY)
+    offenses = inspect_source_within_migration(<<~RUBY)
       add_unique_constraint %i[col_1 col_2], where: "state != 'deleted'"
     RUBY