From 14ace72f7de5a2cec5d860f1ee9bcd024b7c12fc Mon Sep 17 00:00:00 2001
From: Mitsuru Haga <mitsuru-haga@cookpad.com>
Date: Tue, 24 May 2016 17:29:01 +0900
Subject: [PATCH] Memorize original method

---
 lib/stateful_enum/machine.rb                       |  9 +++++----
 test/dummy/app/models/post.rb                      | 13 +++++++++++++
 test/dummy/app/models/quote_post.rb                |  4 ++++
 .../db/migrate/20160524075858_create_posts.rb      | 14 ++++++++++++++
 test/mechanic_machine_test.rb                      | 10 ++++++++++
 5 files changed, 46 insertions(+), 4 deletions(-)
 create mode 100644 test/dummy/app/models/post.rb
 create mode 100644 test/dummy/app/models/quote_post.rb
 create mode 100644 test/dummy/db/migrate/20160524075858_create_posts.rb

diff --git a/lib/stateful_enum/machine.rb b/lib/stateful_enum/machine.rb
index 485ae23..62f7751 100644
--- a/lib/stateful_enum/machine.rb
+++ b/lib/stateful_enum/machine.rb
@@ -3,7 +3,7 @@
 module StatefulEnum
   class Machine
     def initialize(model, column, states, prefix, suffix, &block)
-      @model, @column, @states, @event_names = model, column, states, []
+      @model, @column, @states, @event_names, @original_methods = model, column, states, [], {}
       @prefix = if prefix == true
                   "#{column}_"
                 elsif prefix
@@ -17,6 +17,7 @@ def initialize(model, column, states, prefix, suffix, &block)
 
       # undef non-verb methods e.g. Model#active!
       states.each do |state|
+        @original_methods[state] = @model.instance_method "#{@prefix}#{state}#{@suffix}!"
         @model.send :undef_method, "#{@prefix}#{state}#{@suffix}!"
       end
 
@@ -25,12 +26,12 @@ def initialize(model, column, states, prefix, suffix, &block)
 
     def event(name, &block)
       raise ArgumentError, "event: :#{name} has already been defined." if @event_names.include? name
-      Event.new @model, @column, @states, @prefix, @suffix, name, &block
+      Event.new @model, @column, @states, @original_methods, @prefix, @suffix, name, &block
       @event_names << name
     end
 
     class Event
-      def initialize(model, column, states, prefix, suffix, name, &block)
+      def initialize(model, column, states, original_methods, prefix, suffix, name, &block)
         @states, @name, @transitions, @before, @after = states, name, {}, nil, nil
 
         instance_eval(&block) if block
@@ -48,7 +49,7 @@ def initialize(model, column, states, prefix, suffix, name, &block)
             if to && (!condition || instance_exec(&condition))
               #TODO transaction?
               instance_eval(&before) if before
-              original_method = self.class.send(:_enum_methods_module).instance_method "#{prefix}#{to}#{suffix}!"
+              original_method = original_methods[to]
               ret = original_method.bind(self).call
               instance_eval(&after) if after
               ret
diff --git a/test/dummy/app/models/post.rb b/test/dummy/app/models/post.rb
new file mode 100644
index 0000000..dcafdfb
--- /dev/null
+++ b/test/dummy/app/models/post.rb
@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+class Post < ActiveRecord::Base
+  enum status: {drafted: 0, published: 1} do
+    event :draft do
+      transition :published => :drafted
+    end
+
+    event :publish do
+      transition :drafted => :published
+    end
+  end
+end
diff --git a/test/dummy/app/models/quote_post.rb b/test/dummy/app/models/quote_post.rb
new file mode 100644
index 0000000..645e069
--- /dev/null
+++ b/test/dummy/app/models/quote_post.rb
@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+class QuotePost < Post
+end
diff --git a/test/dummy/db/migrate/20160524075858_create_posts.rb b/test/dummy/db/migrate/20160524075858_create_posts.rb
new file mode 100644
index 0000000..4fc62fc
--- /dev/null
+++ b/test/dummy/db/migrate/20160524075858_create_posts.rb
@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+class CreatePosts < (Rails::VERSION::STRING >= '5' ? ActiveRecord::Migration[5.0] : ActiveRecord::Migration)
+  def change
+    create_table :posts do |t|
+      t.string :title
+      t.text :body
+      t.string :type
+      t.integer :status, default: 0
+
+      t.timestamps null: false
+    end
+  end
+end
diff --git a/test/mechanic_machine_test.rb b/test/mechanic_machine_test.rb
index 8054b90..16f5b79 100644
--- a/test/mechanic_machine_test.rb
+++ b/test/mechanic_machine_test.rb
@@ -104,6 +104,16 @@ def test_unless_condition
     end
   end
 
+  def test_inherited_model
+    post = Post.new
+    post.publish!
+    assert_equal 'published', post.status
+
+    quote_post = QuotePost.new
+    quote_post.publish!
+    assert_equal 'published', quote_post.status
+  end
+
   def test_enum_definition_with_array
     ActiveRecord::Migration.create_table(:array_enum_test) {|t| t.integer :col }
     tes = Class.new(ActiveRecord::Base) do