From c58bbc84134212db756c640d55c097d46d1ab832 Mon Sep 17 00:00:00 2001
From: Yuichi Tanikawa <u1.tanikawa@gmail.com>
Date: Wed, 5 Dec 2018 23:41:32 +0900
Subject: [PATCH] Refactor hooking ActiveRecord migration tasks

- Use Rake::Task#enhance insteaad of defining same tasks again

- Remove hooking db:migrate:change task which doesn't exist

- Fix hooking db:migrate:reset task so that the annotation runs after
  all migration tasks (#548)
---
 lib/tasks/annotate_models_migrate.rake     | 15 +----
 spec/tasks/annotate_models_migrate_spec.rb | 72 ++++++++++++++++++++++
 2 files changed, 75 insertions(+), 12 deletions(-)
 create mode 100644 spec/tasks/annotate_models_migrate_spec.rb

diff --git a/lib/tasks/annotate_models_migrate.rake b/lib/tasks/annotate_models_migrate.rake
index f5a259b22..c8b805c4d 100644
--- a/lib/tasks/annotate_models_migrate.rake
+++ b/lib/tasks/annotate_models_migrate.rake
@@ -4,21 +4,12 @@
 # Append annotations to Rake tasks for ActiveRecord, so annotate automatically gets
 # run after doing db:migrate.
 
-namespace :db do
-  [:migrate, :rollback].each do |cmd|
-    task cmd do
+%w(db:migrate db:migrate:up db:migrate:down db:migrate:reset db:migrate:redo db:rollback).each do |task|
+  Rake::Task[task].enhance do
+    Rake::Task[Rake.application.top_level_tasks.last].enhance do
       Rake::Task['set_annotation_options'].invoke
       Annotate::Migration.update_annotations
     end
-
-    namespace cmd do
-      [:change, :up, :down, :reset, :redo].each do |t|
-        task t do
-          Rake::Task['set_annotation_options'].invoke
-          Annotate::Migration.update_annotations
-        end
-      end
-    end
   end
 end
 
diff --git a/spec/tasks/annotate_models_migrate_spec.rb b/spec/tasks/annotate_models_migrate_spec.rb
new file mode 100644
index 000000000..8c30e1c93
--- /dev/null
+++ b/spec/tasks/annotate_models_migrate_spec.rb
@@ -0,0 +1,72 @@
+require_relative '../spec_helper'
+
+describe 'ActiveRecord migration rake task hooks' do
+  before do
+    Rake.application = Rake::Application.new
+
+    # Stub migration tasks
+    %w(db:migrate db:migrate:up db:migrate:down db:migrate:reset db:rollback).each do |task|
+      Rake::Task.define_task(task)
+    end
+    Rake::Task.define_task('db:migrate:redo') do
+      Rake::Task['db:rollback'].invoke
+      Rake::Task['db:migrate'].invoke
+    end
+
+    Rake::Task.define_task('set_annotation_options')
+    Rake.load_rakefile('tasks/annotate_models_migrate.rake')
+
+    Rake.application.instance_variable_set(:@top_level_tasks, [subject])
+  end
+
+  describe 'db:migrate' do
+    it 'should update annotations' do
+      expect(Annotate::Migration).to receive(:update_annotations)
+      Rake.application.top_level
+    end
+  end
+
+  describe 'db:migrate:up' do
+    it 'should update annotations' do
+      expect(Annotate::Migration).to receive(:update_annotations)
+      Rake.application.top_level
+    end
+  end
+
+  describe 'db:migrate:down' do
+    it 'should update annotations' do
+      expect(Annotate::Migration).to receive(:update_annotations)
+      Rake.application.top_level
+    end
+  end
+
+  describe 'db:migrate:reset' do
+    it 'should update annotations' do
+      expect(Annotate::Migration).to receive(:update_annotations)
+      Rake.application.top_level
+    end
+  end
+
+  describe 'db:rollback' do
+    it 'should update annotations' do
+      expect(Annotate::Migration).to receive(:update_annotations)
+      Rake.application.top_level
+    end
+  end
+
+  describe 'db:migrate:redo' do
+    it 'should update annotations after all migration tasks' do
+      allow(Annotate::Migration).to receive(:update_annotations)
+
+      # Confirm that update_annotations isn't called when the original redo task finishes
+      Rake::Task[subject].enhance do
+        expect(Annotate::Migration).not_to have_received(:update_annotations)
+      end
+
+      Rake.application.top_level
+
+      # Hooked 3 times by db:rollback, db:migrate, and db:migrate:redo tasks
+      expect(Annotate::Migration).to have_received(:update_annotations).exactly(3).times
+    end
+  end
+end