From 8ad414b2175e71a9a00381537cea0ce1cb7813ea Mon Sep 17 00:00:00 2001
From: Sean Hill <seandhi@gmail.com>
Date: Mon, 29 Aug 2011 17:53:32 -0500
Subject: [PATCH] Added a helper to look for an modifying migration.  If one is
 found during :revoke, the modifying migration is deleted.  If the modifying
 migration is not found, the creating migration is deleted.

---
 lib/generators/active_record/devise_generator.rb | 16 ++++++++++++----
 lib/generators/devise/orm_helpers.rb             |  8 ++++++++
 test/generators/active_record_generator_test.rb  |  6 ++++++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/lib/generators/active_record/devise_generator.rb b/lib/generators/active_record/devise_generator.rb
index cf96619d..9d2bd742 100644
--- a/lib/generators/active_record/devise_generator.rb
+++ b/lib/generators/active_record/devise_generator.rb
@@ -1,6 +1,7 @@
 require 'rails/generators/active_record'
 require 'generators/devise/orm_helpers'
 
+
 module ActiveRecord
   module Generators
     class DeviseGenerator < ActiveRecord::Generators::Base
@@ -11,11 +12,18 @@ module ActiveRecord
 
       def copy_devise_migration
         exists = model_exists?
-        exists = !exists if behavior == :revoke
-        unless exists
-          migration_template "migration.rb", "db/migrate/devise_create_#{table_name}"
+        unless behavior == :revoke
+          unless exists
+            migration_template "migration.rb", "db/migrate/devise_create_#{table_name}"
+          else
+            migration_template "migration_existing.rb", "db/migrate/add_devise_to_#{table_name}"
+          end
         else
-          migration_template "migration_existing.rb", "db/migrate/add_devise_to_#{table_name}"
+          if migration_exists?(table_name)
+            migration_template "migration_existing.rb", "db/migrate/add_devise_to_#{table_name}"
+          else
+            migration_template "migration.rb", "db/migrate/devise_create_#{table_name}"
+          end
         end
       end
 
diff --git a/lib/generators/devise/orm_helpers.rb b/lib/generators/devise/orm_helpers.rb
index 99251f28..644f83d1 100644
--- a/lib/generators/devise/orm_helpers.rb
+++ b/lib/generators/devise/orm_helpers.rb
@@ -14,6 +14,14 @@ CONTENT
       def model_exists?
         File.exists?(File.join(destination_root, model_path))
       end
+      
+      def migration_exists?(table_name)
+        Dir.glob("#{File.join(destination_root, migration_path)}/[0-9]*_*.rb").grep(/\d+_add_devise_to_#{table_name}.rb$/).first
+      end
+      
+      def migration_path
+        @migration_path ||= File.join("db", "migrate")
+      end
 
       def model_path
         @model_path ||= File.join("app", "models", "#{file_path}.rb")
diff --git a/test/generators/active_record_generator_test.rb b/test/generators/active_record_generator_test.rb
index 9963e3b8..f2b466ed 100644
--- a/test/generators/active_record_generator_test.rb
+++ b/test/generators/active_record_generator_test.rb
@@ -23,6 +23,12 @@ if DEVISE_ORM == :active_record
   
     test "all files are properly deleted" do
       run_generator %w(monster)
+      run_generator %w(monster)
+      assert_migration "db/migrate/devise_create_monsters.rb"
+      assert_migration "db/migrate/add_devise_to_monsters.rb"
+      run_generator %w(monster), :behavior => :revoke
+      assert_no_migration "db/migrate/add_devise_to_monsters.rb"
+      assert_migration "db/migrate/devise_create_monsters.rb"
       run_generator %w(monster), :behavior => :revoke
       assert_no_file "app/models/monster.rb"
       assert_no_migration "db/migrate/devise_create_monsters.rb"