From 7ad5c748f9861bef908c38bcb6ebc31d6408bf3e Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Tue, 26 Nov 2024 01:26:10 +0800
Subject: [PATCH 1/9] keep enum value state

---
 Gemfile.lock                                  |  1 +
 lib/enum_machine/build_attribute.rb           |  2 +-
 lib/enum_machine/driver_active_record.rb      | 12 ++++++----
 lib/enum_machine/driver_simple_class.rb       |  6 ++---
 spec/enum_machine/active_record_enum_spec.rb  | 16 ++++++++++++++
 spec/enum_machine/driver_simple_class_spec.rb | 22 +++++++++++++++++++
 6 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/Gemfile.lock b/Gemfile.lock
index 302d361..38eaca8 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -146,6 +146,7 @@ GEM
 
 PLATFORMS
   arm64-darwin-21
+  arm64-darwin-23
   x86_64-darwin-21
   x86_64-linux
 
diff --git a/lib/enum_machine/build_attribute.rb b/lib/enum_machine/build_attribute.rb
index 539b0bc..9f48399 100644
--- a/lib/enum_machine/build_attribute.rb
+++ b/lib/enum_machine/build_attribute.rb
@@ -6,7 +6,7 @@ module BuildAttribute
     def self.call(enum_values:, i18n_scope:, machine: nil)
       aliases = machine&.instance_variable_get(:@aliases) || {}
 
-      Class.new(String) do
+      Module.new do
         define_method(:machine) { machine } if machine
 
         def inspect
diff --git a/lib/enum_machine/driver_active_record.rb b/lib/enum_machine/driver_active_record.rb
index 3de63d0..30f89ba 100644
--- a/lib/enum_machine/driver_active_record.rb
+++ b/lib/enum_machine/driver_active_record.rb
@@ -14,13 +14,16 @@ def enum_machine(attr, enum_values, i18n_scope: nil, &block)
 
       enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
 
-      enum_value_klass = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
-      enum_value_klass.extend(AttributePersistenceMethods[attr, enum_values])
+      enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
+      enum_attribute_module.extend(AttributePersistenceMethods[attr, enum_values])
 
-      enum_klass.const_set :VALUE_KLASS, enum_value_klass
+      enum_klass.const_set :ATTRIBUTE_MODULE, enum_attribute_module
 
       # Hash.new with default_proc for working with custom values not defined in enum list
-      value_attribute_mapping = Hash.new { |hash, enum_value| hash[enum_value] = enum_klass::VALUE_KLASS.new(enum_value).freeze }
+      value_attribute_mapping =
+        Hash.new do |hash, enum_value|
+          hash[enum_value] = (enum_values.detect { |value| enum_value == value } || enum_value).dup.extend(enum_klass::ATTRIBUTE_MODULE).freeze
+        end
       enum_klass.define_singleton_method(:value_attribute_mapping) { value_attribute_mapping }
 
       if machine.transitions?
@@ -81,6 +84,7 @@ def #{attr}
 
           unless @__enum_value_#{attr} == enum_value
             @__enum_value_#{attr} = self.class::#{enum_const_name}.value_attribute_mapping[enum_value].dup
+            @__enum_value_#{attr}.extend(self.class::#{enum_const_name}::ATTRIBUTE_MODULE)
             @__enum_value_#{attr}.parent = self
             @__enum_value_#{attr}.freeze
           end
diff --git a/lib/enum_machine/driver_simple_class.rb b/lib/enum_machine/driver_simple_class.rb
index 2d81fd9..0a93b45 100644
--- a/lib/enum_machine/driver_simple_class.rb
+++ b/lib/enum_machine/driver_simple_class.rb
@@ -21,10 +21,10 @@ def self.call(args)
               enum_const_name = attr.to_s.upcase
               enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope)
 
-              enum_value_klass = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope)
-              enum_klass.const_set :VALUE_KLASS, enum_value_klass
+              enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope)
+              enum_klass.const_set :ATTRIBUTE_MODULE, enum_attribute_module
 
-              value_attribute_mapping = enum_values.to_h { |enum_value| [enum_value, enum_klass::VALUE_KLASS.new(enum_value).freeze] }
+              value_attribute_mapping = enum_values.to_h { |enum_value| [enum_value, enum_value.dup.extend(enum_klass::ATTRIBUTE_MODULE).freeze] }
 
               define_methods =
                 Module.new do
diff --git a/spec/enum_machine/active_record_enum_spec.rb b/spec/enum_machine/active_record_enum_spec.rb
index fd2daa8..4483dc1 100644
--- a/spec/enum_machine/active_record_enum_spec.rb
+++ b/spec/enum_machine/active_record_enum_spec.rb
@@ -118,4 +118,20 @@
     expect(decorated_klass::STATE::CHOICE).to eq "choice"
     expect(decorated_klass::COLOR::RED).to eq "red"
   end
+
+  it "keeps class of enum value" do
+    choice_klass =
+      Class.new(String) do
+        def in_choice? = true
+      end
+
+    model =
+      Class.new(TestModel) do
+        enum_machine :state, [choice_klass.new("choice"), "in_delivery"]
+      end
+
+    m = model.new(state: "choice")
+    expect(m.state).to be_a(choice_klass)
+    expect(m.state.in_choice?).to be(true)
+  end
 end
diff --git a/spec/enum_machine/driver_simple_class_spec.rb b/spec/enum_machine/driver_simple_class_spec.rb
index dfd13cb..a087f36 100644
--- a/spec/enum_machine/driver_simple_class_spec.rb
+++ b/spec/enum_machine/driver_simple_class_spec.rb
@@ -92,4 +92,26 @@ def initialize(state)
     expect(decorated_item.state).to be_choice
     expect(decorated_klass::STATE::CHOICE).to eq "choice"
   end
+
+  it "keeps class of enum value" do
+    choice_klass =
+      Class.new(String) do
+        def in_choice? = true
+      end
+
+    test_class =
+      Class.new do
+        attr_accessor :state
+
+        def initialize(state)
+          @state = state
+        end
+
+        include EnumMachine[state: { enum: [choice_klass.new("choice"), "in_delivery"] }]
+      end
+
+    item = test_class.new("choice")
+    expect(item.state).to be_a(choice_klass)
+    expect(item.state.in_choice?).to be(true)
+  end
 end

From 6eab946412659c007375cf130392a94895f847a7 Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Tue, 26 Nov 2024 13:39:02 +0800
Subject: [PATCH 2/9] feat: value_class

---
 lib/enum_machine/driver_active_record.rb      | 20 ++++++++++---------
 lib/enum_machine/driver_simple_class.rb       | 12 +++++++++--
 spec/enum_machine/active_record_enum_spec.rb  |  5 +++--
 spec/enum_machine/driver_simple_class_spec.rb |  5 +++--
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/lib/enum_machine/driver_active_record.rb b/lib/enum_machine/driver_active_record.rb
index 30f89ba..374c8dc 100644
--- a/lib/enum_machine/driver_active_record.rb
+++ b/lib/enum_machine/driver_active_record.rb
@@ -3,7 +3,7 @@
 module EnumMachine
   module DriverActiveRecord
 
-    def enum_machine(attr, enum_values, i18n_scope: nil, &block)
+    def enum_machine(attr, enum_values, i18n_scope: nil, value_class: Class.new(String), &block)
       klass = self
 
       i18n_scope ||= "#{klass.base_class.to_s.underscore}.#{attr}"
@@ -13,18 +13,21 @@ def enum_machine(attr, enum_values, i18n_scope: nil, &block)
       machine.instance_eval(&block) if block
 
       enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
-
       enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
-      enum_attribute_module.extend(AttributePersistenceMethods[attr, enum_values])
 
-      enum_klass.const_set :ATTRIBUTE_MODULE, enum_attribute_module
+      value_class.include(enum_attribute_module)
+      value_class.extend(AttributePersistenceMethods[attr, enum_values])
+
+      enum_klass.const_set(:VALUE_CLASS, value_class)
 
-      # Hash.new with default_proc for working with custom values not defined in enum list
-      value_attribute_mapping =
+      enum_klass.define_singleton_method(:value_attribute_mapping) do
+        # Hash.new with default_proc for working with custom values not defined in enum list
         Hash.new do |hash, enum_value|
-          hash[enum_value] = (enum_values.detect { |value| enum_value == value } || enum_value).dup.extend(enum_klass::ATTRIBUTE_MODULE).freeze
+          value = enum_values.detect { enum_value == _1 } || enum_value
+          value = enum_klass::VALUE_CLASS.new(value) unless value.is_a?(enum_klass::VALUE_CLASS)
+          hash[enum_value] = value.freeze
         end
-      enum_klass.define_singleton_method(:value_attribute_mapping) { value_attribute_mapping }
+      end
 
       if machine.transitions?
         klass.class_eval <<-RUBY, __FILE__, __LINE__ + 1 # rubocop:disable Style/DocumentDynamicEvalDefinition
@@ -84,7 +87,6 @@ def #{attr}
 
           unless @__enum_value_#{attr} == enum_value
             @__enum_value_#{attr} = self.class::#{enum_const_name}.value_attribute_mapping[enum_value].dup
-            @__enum_value_#{attr}.extend(self.class::#{enum_const_name}::ATTRIBUTE_MODULE)
             @__enum_value_#{attr}.parent = self
             @__enum_value_#{attr}.freeze
           end
diff --git a/lib/enum_machine/driver_simple_class.rb b/lib/enum_machine/driver_simple_class.rb
index 0a93b45..1b3fe45 100644
--- a/lib/enum_machine/driver_simple_class.rb
+++ b/lib/enum_machine/driver_simple_class.rb
@@ -14,6 +14,7 @@ def self.call(args)
           args.each do |attr, params|
             enum_values  = params.fetch(:enum)
             i18n_scope   = params.fetch(:i18n_scope, nil)
+            value_class  = params.fetch(:value_class, Class.new(String))
 
             if defined?(ActiveRecord) && klass <= ActiveRecord::Base
               klass.enum_machine(attr, enum_values, i18n_scope: i18n_scope)
@@ -22,9 +23,16 @@ def self.call(args)
               enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope)
 
               enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope)
-              enum_klass.const_set :ATTRIBUTE_MODULE, enum_attribute_module
 
-              value_attribute_mapping = enum_values.to_h { |enum_value| [enum_value, enum_value.dup.extend(enum_klass::ATTRIBUTE_MODULE).freeze] }
+              value_class.include(enum_attribute_module)
+              enum_klass.const_set(:VALUE_CLASS, value_class)
+
+              value_attribute_mapping =
+                enum_values.to_h do |enum_value|
+                  value = enum_values.detect { enum_value == _1 } || enum_value
+                  value = enum_klass::VALUE_CLASS.new(value) unless value.is_a?(enum_klass::VALUE_CLASS)
+                  [enum_value, value.freeze]
+                end
 
               define_methods =
                 Module.new do
diff --git a/spec/enum_machine/active_record_enum_spec.rb b/spec/enum_machine/active_record_enum_spec.rb
index 4483dc1..6e1a51e 100644
--- a/spec/enum_machine/active_record_enum_spec.rb
+++ b/spec/enum_machine/active_record_enum_spec.rb
@@ -120,14 +120,15 @@
   end
 
   it "keeps class of enum value" do
+    value_class = Class.new(String)
     choice_klass =
-      Class.new(String) do
+      Class.new(value_class) do
         def in_choice? = true
       end
 
     model =
       Class.new(TestModel) do
-        enum_machine :state, [choice_klass.new("choice"), "in_delivery"]
+        enum_machine :state, [choice_klass.new("choice"), "in_delivery"], value_class: value_class
       end
 
     m = model.new(state: "choice")
diff --git a/spec/enum_machine/driver_simple_class_spec.rb b/spec/enum_machine/driver_simple_class_spec.rb
index a087f36..11fafbc 100644
--- a/spec/enum_machine/driver_simple_class_spec.rb
+++ b/spec/enum_machine/driver_simple_class_spec.rb
@@ -94,8 +94,9 @@ def initialize(state)
   end
 
   it "keeps class of enum value" do
+    value_class = Class.new(String)
     choice_klass =
-      Class.new(String) do
+      Class.new(value_class) do
         def in_choice? = true
       end
 
@@ -107,7 +108,7 @@ def initialize(state)
           @state = state
         end
 
-        include EnumMachine[state: { enum: [choice_klass.new("choice"), "in_delivery"] }]
+        include EnumMachine[state: { enum: [choice_klass.new("choice"), "in_delivery"], value_class: value_class }]
       end
 
     item = test_class.new("choice")

From 066d1b6752d984a80db1fb87a693a2d5d3b6a44d Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Tue, 26 Nov 2024 13:53:01 +0800
Subject: [PATCH 3/9] fix: value_attribute_mapping

---
 lib/enum_machine/driver_active_record.rb | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/lib/enum_machine/driver_active_record.rb b/lib/enum_machine/driver_active_record.rb
index 374c8dc..4e49214 100644
--- a/lib/enum_machine/driver_active_record.rb
+++ b/lib/enum_machine/driver_active_record.rb
@@ -20,14 +20,14 @@ def enum_machine(attr, enum_values, i18n_scope: nil, value_class: Class.new(Stri
 
       enum_klass.const_set(:VALUE_CLASS, value_class)
 
-      enum_klass.define_singleton_method(:value_attribute_mapping) do
-        # Hash.new with default_proc for working with custom values not defined in enum list
+      # Hash.new with default_proc for working with custom values not defined in enum list
+      value_attribute_mapping =
         Hash.new do |hash, enum_value|
           value = enum_values.detect { enum_value == _1 } || enum_value
           value = enum_klass::VALUE_CLASS.new(value) unless value.is_a?(enum_klass::VALUE_CLASS)
           hash[enum_value] = value.freeze
         end
-      end
+      enum_klass.define_singleton_method(:value_attribute_mapping) { value_attribute_mapping }
 
       if machine.transitions?
         klass.class_eval <<-RUBY, __FILE__, __LINE__ + 1 # rubocop:disable Style/DocumentDynamicEvalDefinition

From f2a7a0788a37d00d4d2810eecf6047a5b4fba5f0 Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Tue, 26 Nov 2024 15:17:24 +0800
Subject: [PATCH 4/9] feat: get state value by `[]`

---
 lib/enum_machine/build_class.rb               | 14 +++++++++++---
 lib/enum_machine/driver_active_record.rb      | 13 +++++--------
 lib/enum_machine/driver_simple_class.rb       | 12 ++----------
 spec/enum_machine/active_record_enum_spec.rb  |  7 +++++++
 spec/enum_machine/driver_simple_class_spec.rb |  7 +++++++
 5 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/lib/enum_machine/build_class.rb b/lib/enum_machine/build_class.rb
index b3b8129..c5cb88e 100644
--- a/lib/enum_machine/build_class.rb
+++ b/lib/enum_machine/build_class.rb
@@ -3,12 +3,20 @@
 module EnumMachine
   module BuildClass
 
-    def self.call(enum_values:, i18n_scope:, machine: nil)
+    def self.call(enum_values:, i18n_scope:, value_class:, machine: nil)
       aliases = machine&.instance_variable_get(:@aliases) || {}
 
       Class.new do
         define_singleton_method(:machine) { machine } if machine
-        define_singleton_method(:values) { enum_values }
+        define_singleton_method(:values) { enum_values.map { _1.is_a?(value_class) ? _1.freeze : value_class.new(_1).freeze } }
+
+        value_attribute_mapping = values.to_h { [_1.to_s, _1] }
+
+        define_singleton_method(:value_attribute_mapping) { value_attribute_mapping }
+        define_singleton_method(:[]) do |enum_value|
+          key = enum_value.to_s
+          value_attribute_mapping.fetch(key) if value_attribute_mapping.key?(key)
+        end
 
         if i18n_scope
           def self.values_for_form(specific_values = nil) # rubocop:disable Gp/OptArgParameters
@@ -27,7 +35,7 @@ def self.human_name_for(name)
         end
 
         enum_values.each do |enum_value|
-          const_set enum_value.underscore.upcase, enum_value.freeze
+          const_set enum_value.underscore.upcase, enum_value.to_s.freeze
         end
 
         aliases.each_key do |key|
diff --git a/lib/enum_machine/driver_active_record.rb b/lib/enum_machine/driver_active_record.rb
index 4e49214..a8a87a0 100644
--- a/lib/enum_machine/driver_active_record.rb
+++ b/lib/enum_machine/driver_active_record.rb
@@ -12,7 +12,7 @@ def enum_machine(attr, enum_values, i18n_scope: nil, value_class: Class.new(Stri
       machine = Machine.new(enum_values, klass, enum_const_name, attr)
       machine.instance_eval(&block) if block
 
-      enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
+      enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine, value_class: value_class)
       enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
 
       value_class.include(enum_attribute_module)
@@ -20,14 +20,11 @@ def enum_machine(attr, enum_values, i18n_scope: nil, value_class: Class.new(Stri
 
       enum_klass.const_set(:VALUE_CLASS, value_class)
 
-      # Hash.new with default_proc for working with custom values not defined in enum list
-      value_attribute_mapping =
-        Hash.new do |hash, enum_value|
-          value = enum_values.detect { enum_value == _1 } || enum_value
-          value = enum_klass::VALUE_CLASS.new(value) unless value.is_a?(enum_klass::VALUE_CLASS)
-          hash[enum_value] = value.freeze
+      # default_proc for working with custom values not defined in enum list but may exists in db
+      enum_klass.value_attribute_mapping.default_proc =
+        proc do |hash, enum_value|
+          hash[enum_value] = value_class.new(enum_value).freeze
         end
-      enum_klass.define_singleton_method(:value_attribute_mapping) { value_attribute_mapping }
 
       if machine.transitions?
         klass.class_eval <<-RUBY, __FILE__, __LINE__ + 1 # rubocop:disable Style/DocumentDynamicEvalDefinition
diff --git a/lib/enum_machine/driver_simple_class.rb b/lib/enum_machine/driver_simple_class.rb
index 1b3fe45..003b2a2 100644
--- a/lib/enum_machine/driver_simple_class.rb
+++ b/lib/enum_machine/driver_simple_class.rb
@@ -20,27 +20,19 @@ def self.call(args)
               klass.enum_machine(attr, enum_values, i18n_scope: i18n_scope)
             else
               enum_const_name = attr.to_s.upcase
-              enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope)
-
+              enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope, value_class: value_class)
               enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope)
 
               value_class.include(enum_attribute_module)
               enum_klass.const_set(:VALUE_CLASS, value_class)
 
-              value_attribute_mapping =
-                enum_values.to_h do |enum_value|
-                  value = enum_values.detect { enum_value == _1 } || enum_value
-                  value = enum_klass::VALUE_CLASS.new(value) unless value.is_a?(enum_klass::VALUE_CLASS)
-                  [enum_value, value.freeze]
-                end
-
               define_methods =
                 Module.new do
                   define_method(attr) do
                     enum_value = super()
                     return unless enum_value
 
-                    value_attribute_mapping.fetch(enum_value)
+                    enum_klass.value_attribute_mapping.fetch(enum_value)
                   end
                 end
 
diff --git a/spec/enum_machine/active_record_enum_spec.rb b/spec/enum_machine/active_record_enum_spec.rb
index 6e1a51e..4292fba 100644
--- a/spec/enum_machine/active_record_enum_spec.rb
+++ b/spec/enum_machine/active_record_enum_spec.rb
@@ -135,4 +135,11 @@ def in_choice? = true
     expect(m.state).to be_a(choice_klass)
     expect(m.state.in_choice?).to be(true)
   end
+
+  it "returns state value by []" do
+    expect(model::STATE["in_delivery"]).to eq "in_delivery"
+    expect(model::STATE["in_delivery"].in_delivery?).to be(true)
+    expect(model::STATE["in_delivery"].choice?).to be(false)
+    expect(model::STATE["wrong"]).to be_nil
+  end
 end
diff --git a/spec/enum_machine/driver_simple_class_spec.rb b/spec/enum_machine/driver_simple_class_spec.rb
index 11fafbc..eee3451 100644
--- a/spec/enum_machine/driver_simple_class_spec.rb
+++ b/spec/enum_machine/driver_simple_class_spec.rb
@@ -115,4 +115,11 @@ def initialize(state)
     expect(item.state).to be_a(choice_klass)
     expect(item.state.in_choice?).to be(true)
   end
+
+  it "returns state value by []" do
+    expect(TestClass::STATE["in_delivery"]).to eq "in_delivery"
+    expect(TestClass::STATE["in_delivery"].in_delivery?).to be(true)
+    expect(TestClass::STATE["in_delivery"].choice?).to be(false)
+    expect(TestClass::STATE["wrong"]).to be_nil
+  end
 end

From 9ed10decf931603c8331a342c445dcb88ec6f248 Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Tue, 26 Nov 2024 17:46:48 +0800
Subject: [PATCH 5/9] chore: readme

---
 README.md                       | 42 +++++++++++++++++++++++++++++----
 lib/enum_machine/build_class.rb |  4 ++--
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 1af78a1..7d88586 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
-# Enum machine
+# Enum Machine
 
-Enum machine is a library for defining enums and setting state machines for attributes in ActiveRecord models and plain Ruby classes.
+`Enum Machine` is a library for defining enums and setting state machines for attributes in ActiveRecord models and plain Ruby classes.
 
 You can visualize transitions map with [enum_machine-contrib](https://github.com/corp-gp/enum_machine-contrib)
 
@@ -29,7 +29,7 @@ end
 
 order = Order.create(state: "collecting")
 order.update(state: "archived") # not check transitions, invalid logic
-order.update(state: "collected") # not run callbacks 
+order.update(state: "collected") # not run callbacks
 order.complete # need use event for transition, but your object in UI and DB have only states
 
 # enum_machine
@@ -66,7 +66,7 @@ end
 class Product
   # attributes must be defined before including the EnumMachine module
   attr_accessor :color
-  
+
   include EnumMachine[color: { enum: %w[red green] }]
   # or reuse from model
   Product::COLOR.decorator_module
@@ -76,6 +76,9 @@ Product::COLOR.values # => ["red", "green"]
 Product::COLOR::RED # => "red"
 Product::COLOR::RED__GREEN # => ["red", "green"]
 
+Product::COLOR["red"].red? # => true
+Product::COLOR["red"].human_name # => "Красный"
+
 product = Product.new
 product.color # => nil
 product.color = "red"
@@ -100,6 +103,37 @@ product = Product.new(state: "created")
 product.state.forming? # => true
 ```
 
+### Custom value classes
+
+You can use custom classes as enum values instead of strings.
+
+```ruby
+# Base value class nested from String
+class EnumValue < String; end
+
+# Value classes nested from base class
+class EnumValueRed < EnumValue
+  def hex = "#ff0000"
+end
+
+class EnumValueGreen < EnumValue
+  def hex = "#00ff00"
+end
+
+class Product
+  attr_accessor :color
+
+  include EnumMachine[color: {
+    enum:        [EnumValueRed.new("red"), EnumValueGreen.new("green")],
+    value_class: EnumValue # Enum Machine will extend this class with it's own methods
+  }]
+end
+
+product = Product.new
+product.color = "red"
+product.color.hex # => "#ff0000"
+```
+
 ### Transitions
 
 ```ruby
diff --git a/lib/enum_machine/build_class.rb b/lib/enum_machine/build_class.rb
index c5cb88e..ed9335f 100644
--- a/lib/enum_machine/build_class.rb
+++ b/lib/enum_machine/build_class.rb
@@ -11,11 +11,11 @@ def self.call(enum_values:, i18n_scope:, value_class:, machine: nil)
         define_singleton_method(:values) { enum_values.map { _1.is_a?(value_class) ? _1.freeze : value_class.new(_1).freeze } }
 
         value_attribute_mapping = values.to_h { [_1.to_s, _1] }
-
         define_singleton_method(:value_attribute_mapping) { value_attribute_mapping }
         define_singleton_method(:[]) do |enum_value|
           key = enum_value.to_s
-          value_attribute_mapping.fetch(key) if value_attribute_mapping.key?(key)
+          # Check for key existence because `[]` will call `default_proc`, and we don’t want that
+          value_attribute_mapping[key] if value_attribute_mapping.key?(key)
         end
 
         if i18n_scope

From 3aa4c2fee9c4a7bb9fdb9f17fb05f992d12b15f7 Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Tue, 26 Nov 2024 18:35:09 +0800
Subject: [PATCH 6/9] feat: decorator

---
 README.md                                     | 24 +++++++++----------
 lib/enum_machine/build_attribute.rb           |  6 +++--
 lib/enum_machine/build_class.rb               |  2 +-
 lib/enum_machine/driver_active_record.rb      |  7 ++----
 lib/enum_machine/driver_simple_class.rb       |  7 ++----
 spec/enum_machine/active_record_enum_spec.rb  | 16 ++++++-------
 spec/enum_machine/driver_simple_class_spec.rb | 16 ++++++-------
 7 files changed, 36 insertions(+), 42 deletions(-)

diff --git a/README.md b/README.md
index 7d88586..38583ad 100644
--- a/README.md
+++ b/README.md
@@ -103,29 +103,27 @@ product = Product.new(state: "created")
 product.state.forming? # => true
 ```
 
-### Custom value classes
+### Value decorator
 
-You can use custom classes as enum values instead of strings.
+You can extend value object with decorator
 
 ```ruby
-# Base value class nested from String
-class EnumValue < String; end
-
 # Value classes nested from base class
-class EnumValueRed < EnumValue
-  def hex = "#ff0000"
-end
-
-class EnumValueGreen < EnumValue
-  def hex = "#00ff00"
+module ColorDecorator
+  def hex
+    case self
+    when Product::COLOR::RED then "#ff0000"
+    when Product::COLOR::GREEN then "#00ff00"
+    end
+  end
 end
 
 class Product
   attr_accessor :color
 
   include EnumMachine[color: {
-    enum:        [EnumValueRed.new("red"), EnumValueGreen.new("green")],
-    value_class: EnumValue # Enum Machine will extend this class with it's own methods
+    enum:      %w[red green],
+    decorator: ColorDecorator
   }]
 end
 
diff --git a/lib/enum_machine/build_attribute.rb b/lib/enum_machine/build_attribute.rb
index 9f48399..c745fe6 100644
--- a/lib/enum_machine/build_attribute.rb
+++ b/lib/enum_machine/build_attribute.rb
@@ -3,10 +3,12 @@
 module EnumMachine
   module BuildAttribute
 
-    def self.call(enum_values:, i18n_scope:, machine: nil)
+    def self.call(enum_values:, i18n_scope:, decorator:, machine: nil)
       aliases = machine&.instance_variable_get(:@aliases) || {}
 
-      Module.new do
+      Class.new(String) do
+        include(decorator) if decorator
+
         define_method(:machine) { machine } if machine
 
         def inspect
diff --git a/lib/enum_machine/build_class.rb b/lib/enum_machine/build_class.rb
index ed9335f..23ab2a1 100644
--- a/lib/enum_machine/build_class.rb
+++ b/lib/enum_machine/build_class.rb
@@ -8,7 +8,7 @@ def self.call(enum_values:, i18n_scope:, value_class:, machine: nil)
 
       Class.new do
         define_singleton_method(:machine) { machine } if machine
-        define_singleton_method(:values) { enum_values.map { _1.is_a?(value_class) ? _1.freeze : value_class.new(_1).freeze } }
+        define_singleton_method(:values) { enum_values.map { value_class.new(_1).freeze } }
 
         value_attribute_mapping = values.to_h { [_1.to_s, _1] }
         define_singleton_method(:value_attribute_mapping) { value_attribute_mapping }
diff --git a/lib/enum_machine/driver_active_record.rb b/lib/enum_machine/driver_active_record.rb
index a8a87a0..2da3525 100644
--- a/lib/enum_machine/driver_active_record.rb
+++ b/lib/enum_machine/driver_active_record.rb
@@ -3,7 +3,7 @@
 module EnumMachine
   module DriverActiveRecord
 
-    def enum_machine(attr, enum_values, i18n_scope: nil, value_class: Class.new(String), &block)
+    def enum_machine(attr, enum_values, i18n_scope: nil, decorator: nil, &block)
       klass = self
 
       i18n_scope ||= "#{klass.base_class.to_s.underscore}.#{attr}"
@@ -12,14 +12,11 @@ def enum_machine(attr, enum_values, i18n_scope: nil, value_class: Class.new(Stri
       machine = Machine.new(enum_values, klass, enum_const_name, attr)
       machine.instance_eval(&block) if block
 
+      value_class = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine, decorator: decorator)
       enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine, value_class: value_class)
-      enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope, machine: machine)
 
-      value_class.include(enum_attribute_module)
       value_class.extend(AttributePersistenceMethods[attr, enum_values])
 
-      enum_klass.const_set(:VALUE_CLASS, value_class)
-
       # default_proc for working with custom values not defined in enum list but may exists in db
       enum_klass.value_attribute_mapping.default_proc =
         proc do |hash, enum_value|
diff --git a/lib/enum_machine/driver_simple_class.rb b/lib/enum_machine/driver_simple_class.rb
index 003b2a2..8a4f0d8 100644
--- a/lib/enum_machine/driver_simple_class.rb
+++ b/lib/enum_machine/driver_simple_class.rb
@@ -14,17 +14,14 @@ def self.call(args)
           args.each do |attr, params|
             enum_values  = params.fetch(:enum)
             i18n_scope   = params.fetch(:i18n_scope, nil)
-            value_class  = params.fetch(:value_class, Class.new(String))
+            decorator    = params.fetch(:decorator, nil)
 
             if defined?(ActiveRecord) && klass <= ActiveRecord::Base
               klass.enum_machine(attr, enum_values, i18n_scope: i18n_scope)
             else
               enum_const_name = attr.to_s.upcase
+              value_class = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope, decorator: decorator)
               enum_klass = BuildClass.call(enum_values: enum_values, i18n_scope: i18n_scope, value_class: value_class)
-              enum_attribute_module = BuildAttribute.call(enum_values: enum_values, i18n_scope: i18n_scope)
-
-              value_class.include(enum_attribute_module)
-              enum_klass.const_set(:VALUE_CLASS, value_class)
 
               define_methods =
                 Module.new do
diff --git a/spec/enum_machine/active_record_enum_spec.rb b/spec/enum_machine/active_record_enum_spec.rb
index 4292fba..aee0182 100644
--- a/spec/enum_machine/active_record_enum_spec.rb
+++ b/spec/enum_machine/active_record_enum_spec.rb
@@ -120,20 +120,20 @@
   end
 
   it "keeps class of enum value" do
-    value_class = Class.new(String)
-    choice_klass =
-      Class.new(value_class) do
-        def in_choice? = true
+    decorator =
+      Module.new do
+        def am_i_choice?
+          self == "choice"
+        end
       end
 
     model =
       Class.new(TestModel) do
-        enum_machine :state, [choice_klass.new("choice"), "in_delivery"], value_class: value_class
+        enum_machine :state, %w[choice in_delivery], decorator: decorator
       end
 
-    m = model.new(state: "choice")
-    expect(m.state).to be_a(choice_klass)
-    expect(m.state.in_choice?).to be(true)
+    expect(model.new(state: "choice").state.am_i_choice?).to be(true)
+    expect(model.new(state: "in_delivery").state.am_i_choice?).to be(false)
   end
 
   it "returns state value by []" do
diff --git a/spec/enum_machine/driver_simple_class_spec.rb b/spec/enum_machine/driver_simple_class_spec.rb
index eee3451..8853d54 100644
--- a/spec/enum_machine/driver_simple_class_spec.rb
+++ b/spec/enum_machine/driver_simple_class_spec.rb
@@ -94,10 +94,11 @@ def initialize(state)
   end
 
   it "keeps class of enum value" do
-    value_class = Class.new(String)
-    choice_klass =
-      Class.new(value_class) do
-        def in_choice? = true
+    decorator =
+      Module.new do
+        def am_i_choice?
+          self == "choice"
+        end
       end
 
     test_class =
@@ -108,12 +109,11 @@ def initialize(state)
           @state = state
         end
 
-        include EnumMachine[state: { enum: [choice_klass.new("choice"), "in_delivery"], value_class: value_class }]
+        include EnumMachine[state: { enum: %w[choice in_delivery], decorator: decorator }]
       end
 
-    item = test_class.new("choice")
-    expect(item.state).to be_a(choice_klass)
-    expect(item.state.in_choice?).to be(true)
+    expect(test_class.new("choice").state.am_i_choice?).to be(true)
+    expect(test_class.new("in_delivery").state.am_i_choice?).to be(false)
   end
 
   it "returns state value by []" do

From 312f55d18a04116f778625bd0668209e1aafffd7 Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Wed, 27 Nov 2024 11:32:14 +0800
Subject: [PATCH 7/9] test: `#values` tests

---
 Gemfile.lock                                  |  30 ++---
 lib/enum_machine.rb                           |   4 -
 .../attribute_persistence_methods.rb          |   2 -
 lib/enum_machine/build_attribute.rb           |   2 -
 lib/enum_machine/build_class.rb               |   2 -
 lib/enum_machine/driver_active_record.rb      |   2 -
 lib/enum_machine/driver_simple_class.rb       |   2 -
 lib/enum_machine/machine.rb                   |   3 -
 lib/enum_machine/version.rb                   |   2 -
 spec/enum_machine/driver_simple_class_spec.rb | 105 ++++++++++--------
 spec/support/test_model.rb                    |   2 -
 test/performance.rb                           |   6 -
 12 files changed, 76 insertions(+), 86 deletions(-)

diff --git a/Gemfile.lock b/Gemfile.lock
index 38eaca8..101aceb 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -8,7 +8,7 @@ GIT
 
 GIT
   remote: https://github.com/corp-gp/rubocop-gp.git
-  revision: 61d6ca243e1a5f7d209e9619417c8fc1d303d181
+  revision: 867f7e1351c3730897cacdc20ce2d727604de245
   specs:
     rubocop-gp (0.0.4)
       rubocop
@@ -73,20 +73,20 @@ GEM
       zeitwerk (~> 2.6)
     i18n (1.14.6)
       concurrent-ruby (~> 1.0)
-    json (2.7.2)
+    json (2.8.2)
     language_server-protocol (3.17.0.3)
     logger (1.6.1)
     method_source (1.1.0)
-    minitest (5.25.1)
+    minitest (5.25.2)
     parallel (1.26.3)
-    parser (3.3.5.0)
+    parser (3.3.6.0)
       ast (~> 2.4.1)
       racc
     pry (0.14.2)
       coderay (~> 1.1)
       method_source (~> 1.0)
     racc (1.8.1)
-    rack (3.1.7)
+    rack (3.1.8)
     rainbow (3.1.1)
     rake (13.2.1)
     regexp_parser (2.9.2)
@@ -103,45 +103,47 @@ GEM
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.13.0)
     rspec-support (3.13.1)
-    rubocop (1.66.1)
+    rubocop (1.69.0)
       json (~> 2.3)
       language_server-protocol (>= 3.17.0)
       parallel (~> 1.10)
       parser (>= 3.3.0.2)
       rainbow (>= 2.2.2, < 4.0)
       regexp_parser (>= 2.4, < 3.0)
-      rubocop-ast (>= 1.32.2, < 2.0)
+      rubocop-ast (>= 1.36.1, < 2.0)
       ruby-progressbar (~> 1.7)
-      unicode-display_width (>= 2.4.0, < 3.0)
-    rubocop-ast (1.32.3)
+      unicode-display_width (>= 2.4.0, < 4.0)
+    rubocop-ast (1.36.1)
       parser (>= 3.3.1.0)
     rubocop-capybara (2.21.0)
       rubocop (~> 1.41)
     rubocop-factory_bot (2.26.1)
       rubocop (~> 1.61)
-    rubocop-performance (1.22.1)
+    rubocop-performance (1.23.0)
       rubocop (>= 1.48.1, < 2.0)
       rubocop-ast (>= 1.31.1, < 2.0)
-    rubocop-rails (2.26.2)
+    rubocop-rails (2.27.0)
       activesupport (>= 4.2.0)
       rack (>= 1.1)
       rubocop (>= 1.52.0, < 2.0)
       rubocop-ast (>= 1.31.1, < 2.0)
-    rubocop-rspec (3.0.5)
+    rubocop-rspec (3.2.0)
       rubocop (~> 1.61)
     rubocop-rspec_rails (2.30.0)
       rubocop (~> 1.61)
       rubocop-rspec (~> 3, >= 3.0.1)
     ruby-progressbar (1.13.0)
     rumoji (0.5.0)
-    securerandom (0.3.1)
+    securerandom (0.3.2)
     sqlite3 (1.7.3-arm64-darwin)
     sqlite3 (1.7.3-x86_64-darwin)
     sqlite3 (1.7.3-x86_64-linux)
     timeout (0.4.1)
     tzinfo (2.0.6)
       concurrent-ruby (~> 1.0)
-    unicode-display_width (2.6.0)
+    unicode-display_width (3.1.2)
+      unicode-emoji (~> 4.0, >= 4.0.4)
+    unicode-emoji (4.0.4)
     zeitwerk (2.6.18)
 
 PLATFORMS
diff --git a/lib/enum_machine.rb b/lib/enum_machine.rb
index 63c32c7..a5db28c 100644
--- a/lib/enum_machine.rb
+++ b/lib/enum_machine.rb
@@ -9,11 +9,9 @@
 require "active_support"
 
 module EnumMachine
-
   class Error < StandardError; end
 
   class InvalidTransition < Error
-
     attr_reader :from, :to, :enum_const
 
     def initialize(machine, from, to)
@@ -26,13 +24,11 @@ def initialize(machine, from, to)
         end
       super("Transition #{from.inspect} => #{to.inspect} not defined in enum_machine :#{machine.attr_name}")
     end
-
   end
 
   def self.[](args)
     DriverSimpleClass.call(args)
   end
-
 end
 
 ActiveSupport.on_load(:active_record) do
diff --git a/lib/enum_machine/attribute_persistence_methods.rb b/lib/enum_machine/attribute_persistence_methods.rb
index e845ba7..5fb8dd0 100644
--- a/lib/enum_machine/attribute_persistence_methods.rb
+++ b/lib/enum_machine/attribute_persistence_methods.rb
@@ -2,7 +2,6 @@
 
 module EnumMachine
   module AttributePersistenceMethods
-
     def self.[](attr, enum_values)
       Module.new do
         define_singleton_method(:extended) do |klass|
@@ -28,6 +27,5 @@ def to_#{enum_name}!
         end
       end
     end
-
   end
 end
diff --git a/lib/enum_machine/build_attribute.rb b/lib/enum_machine/build_attribute.rb
index c745fe6..4a3e9fc 100644
--- a/lib/enum_machine/build_attribute.rb
+++ b/lib/enum_machine/build_attribute.rb
@@ -2,7 +2,6 @@
 
 module EnumMachine
   module BuildAttribute
-
     def self.call(enum_values:, i18n_scope:, decorator:, machine: nil)
       aliases = machine&.instance_variable_get(:@aliases) || {}
 
@@ -76,6 +75,5 @@ def human_name
         end
       end
     end
-
   end
 end
diff --git a/lib/enum_machine/build_class.rb b/lib/enum_machine/build_class.rb
index 23ab2a1..fc2667a 100644
--- a/lib/enum_machine/build_class.rb
+++ b/lib/enum_machine/build_class.rb
@@ -2,7 +2,6 @@
 
 module EnumMachine
   module BuildClass
-
     def self.call(enum_values:, i18n_scope:, value_class:, machine: nil)
       aliases = machine&.instance_variable_get(:@aliases) || {}
 
@@ -58,6 +57,5 @@ def self.#{key}
         end
       end
     end
-
   end
 end
diff --git a/lib/enum_machine/driver_active_record.rb b/lib/enum_machine/driver_active_record.rb
index 2da3525..5368db3 100644
--- a/lib/enum_machine/driver_active_record.rb
+++ b/lib/enum_machine/driver_active_record.rb
@@ -2,7 +2,6 @@
 
 module EnumMachine
   module DriverActiveRecord
-
     def enum_machine(attr, enum_values, i18n_scope: nil, decorator: nil, &block)
       klass = self
 
@@ -114,6 +113,5 @@ def initialize_dup(other)
 
       enum_decorator
     end
-
   end
 end
diff --git a/lib/enum_machine/driver_simple_class.rb b/lib/enum_machine/driver_simple_class.rb
index 8a4f0d8..6a566c5 100644
--- a/lib/enum_machine/driver_simple_class.rb
+++ b/lib/enum_machine/driver_simple_class.rb
@@ -2,7 +2,6 @@
 
 module EnumMachine
   module DriverSimpleClass
-
     # include EnumMachine[
     #   state: { enum: %w[choice in_delivery], i18n_scope: 'line_item.state' },
     #   color: { enum: %w[red green yellow] },
@@ -49,6 +48,5 @@ def self.call(args)
         end
       end
     end
-
   end
 end
diff --git a/lib/enum_machine/machine.rb b/lib/enum_machine/machine.rb
index c7e5118..7e5dea5 100644
--- a/lib/enum_machine/machine.rb
+++ b/lib/enum_machine/machine.rb
@@ -2,7 +2,6 @@
 
 module EnumMachine
   class Machine
-
     attr_reader :enum_values, :base_klass, :enum_const_name, :attr_name
 
     def initialize(enum_values, base_klass = nil, enum_const_name = nil, attr_name = nil) # rubocop:disable Gp/OptArgParameters
@@ -132,8 +131,6 @@ def possible_transitions(from)
     end
 
     class AnyEnumValues < Array
-
     end
-
   end
 end
diff --git a/lib/enum_machine/version.rb b/lib/enum_machine/version.rb
index 8c8bb34..171f209 100644
--- a/lib/enum_machine/version.rb
+++ b/lib/enum_machine/version.rb
@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
 module EnumMachine
-
   VERSION = "1.0.0"
-
 end
diff --git a/spec/enum_machine/driver_simple_class_spec.rb b/spec/enum_machine/driver_simple_class_spec.rb
index 8853d54..66a7253 100644
--- a/spec/enum_machine/driver_simple_class_spec.rb
+++ b/spec/enum_machine/driver_simple_class_spec.rb
@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 class TestClass
-
   attr_accessor :state
 
   def initialize(state)
@@ -9,7 +8,22 @@ def initialize(state)
   end
 
   include EnumMachine[state: { enum: %w[choice in_delivery] }]
+end
 
+module Decorator
+  def am_i_choice?
+    self == "choice"
+  end
+end
+
+class TestClassWithDecorator
+  attr_accessor :state
+
+  def initialize(state)
+    @state = state
+  end
+
+  include EnumMachine[state: { enum: %w[choice in_delivery], decorator: Decorator }]
 end
 
 RSpec.describe "DriverSimpleClass" do
@@ -52,6 +66,33 @@ def initialize(state)
     end
   end
 
+  describe "TestClass::STATE const" do
+    it "#values" do
+      expect(TestClass::STATE.values).to eq(%w[choice in_delivery])
+    end
+
+    it "#[]" do
+      expect(TestClass::STATE["in_delivery"]).to eq "in_delivery"
+      expect(TestClass::STATE["in_delivery"].in_delivery?).to be(true)
+      expect(TestClass::STATE["in_delivery"].choice?).to be(false)
+      expect(TestClass::STATE["wrong"]).to be_nil
+    end
+
+    it "#decorator_module" do
+      decorated_klass =
+        Class.new do
+          include TestClass::STATE.decorator_module
+          attr_accessor :state
+        end
+
+      decorated_item = decorated_klass.new
+      decorated_item.state = "choice"
+
+      expect(decorated_item.state).to be_choice
+      expect(decorated_klass::STATE::CHOICE).to eq "choice"
+    end
+  end
+
   context "when definition order is changed" do
     let(:invert_definition_class) do
       Class.new do
@@ -70,6 +111,24 @@ def initialize(state)
     end
   end
 
+  context "when with decorator" do
+    it "decorates enum values" do
+      expect(TestClassWithDecorator.new("choice").state.am_i_choice?).to be(true)
+      expect(TestClassWithDecorator.new("in_delivery").state.am_i_choice?).to be(false)
+    end
+
+    it "decorates enum values in enum const" do
+      expect(TestClassWithDecorator::STATE.values.map(&:am_i_choice?)).to eq([true, false])
+      expect((TestClassWithDecorator::STATE.values & ["in_delivery"]).map(&:am_i_choice?)).to eq([false])
+    end
+
+    it "keeps decorating on serialization" do
+      m = TestClassWithDecorator.new("choice")
+      unserialized_m = Marshal.load(Marshal.dump(m)) # rubocop:disable Gp/UnsafeYamlMarshal
+      expect(unserialized_m.state.am_i_choice?).to be(true)
+    end
+  end
+
   it "serialize class" do
     m = TestClass.new("choice")
 
@@ -78,48 +137,4 @@ def initialize(state)
     expect(unserialized_m.state).to be_choice
     expect(unserialized_m.class::STATE::CHOICE).to eq "choice"
   end
-
-  it "test decorator" do
-    decorated_klass =
-      Class.new do
-        include TestClass::STATE.decorator_module
-        attr_accessor :state
-      end
-
-    decorated_item = decorated_klass.new
-    decorated_item.state = "choice"
-
-    expect(decorated_item.state).to be_choice
-    expect(decorated_klass::STATE::CHOICE).to eq "choice"
-  end
-
-  it "keeps class of enum value" do
-    decorator =
-      Module.new do
-        def am_i_choice?
-          self == "choice"
-        end
-      end
-
-    test_class =
-      Class.new do
-        attr_accessor :state
-
-        def initialize(state)
-          @state = state
-        end
-
-        include EnumMachine[state: { enum: %w[choice in_delivery], decorator: decorator }]
-      end
-
-    expect(test_class.new("choice").state.am_i_choice?).to be(true)
-    expect(test_class.new("in_delivery").state.am_i_choice?).to be(false)
-  end
-
-  it "returns state value by []" do
-    expect(TestClass::STATE["in_delivery"]).to eq "in_delivery"
-    expect(TestClass::STATE["in_delivery"].in_delivery?).to be(true)
-    expect(TestClass::STATE["in_delivery"].choice?).to be(false)
-    expect(TestClass::STATE["wrong"]).to be_nil
-  end
 end
diff --git a/spec/support/test_model.rb b/spec/support/test_model.rb
index c81d481..0b4a516 100644
--- a/spec/support/test_model.rb
+++ b/spec/support/test_model.rb
@@ -3,9 +3,7 @@
 require "active_record"
 
 class TestModel < ActiveRecord::Base
-
   def self.model_name
     ActiveModel::Name.new(self, nil, "test_model")
   end
-
 end
diff --git a/test/performance.rb b/test/performance.rb
index e842f7f..8206117 100644
--- a/test/performance.rb
+++ b/test/performance.rb
@@ -32,7 +32,6 @@
 STATES_IN_TRANSIT = %w[shipped delivered_to_office delivered_to_courier_city].freeze
 
 class OrderEnumMachine < ActiveRecord::Base
-
   self.table_name = :orders
 
   enum_machine :state, %w[
@@ -62,11 +61,9 @@ class OrderEnumMachine < ActiveRecord::Base
       %w[wait_shipment shipped part_obtain obtain overdue rejection searched] | STATES_IN_TRANSIT => "lost",
     )
   end
-
 end
 
 class OrderAasm < ActiveRecord::Base
-
   include AASM
 
   self.table_name = :orders
@@ -175,11 +172,9 @@ class OrderAasm < ActiveRecord::Base
       transitions from: %i[wait_shipment shipped part_obtain obtain overdue rejection searched] | STATES_IN_TRANSIT, to: :lost
     end
   end
-
 end
 
 class OrderStateMachines < ActiveRecord::Base
-
   self.table_name = :orders
 
   state_machine :state, initial: nil do
@@ -262,7 +257,6 @@ class OrderStateMachines < ActiveRecord::Base
       transition %w[wait_shipment shipped part_obtain obtain overdue rejection searched] | STATES_IN_TRANSIT => "lost"
     end
   end
-
 end
 
 def pp_title(name, stmt)

From ad15370fac9e169bf516306b3dbd32bd10c95b46 Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Wed, 27 Nov 2024 11:57:38 +0800
Subject: [PATCH 8/9] feat: inspect

---
 lib/enum_machine/attribute_persistence_methods.rb | 4 ----
 lib/enum_machine/build_attribute.rb               | 2 +-
 spec/enum_machine/active_record_enum_spec.rb      | 5 -----
 3 files changed, 1 insertion(+), 10 deletions(-)

diff --git a/lib/enum_machine/attribute_persistence_methods.rb b/lib/enum_machine/attribute_persistence_methods.rb
index 5fb8dd0..8d1bfce 100644
--- a/lib/enum_machine/attribute_persistence_methods.rb
+++ b/lib/enum_machine/attribute_persistence_methods.rb
@@ -7,10 +7,6 @@ def self.[](attr, enum_values)
         define_singleton_method(:extended) do |klass|
           klass.attr_accessor :parent
 
-          klass.define_method(:inspect) do
-            "#<EnumMachine:BuildAttribute value=#{self} parent=#{parent.inspect}>"
-          end
-
           enum_values.each do |enum_value|
             enum_name = enum_value.underscore
 
diff --git a/lib/enum_machine/build_attribute.rb b/lib/enum_machine/build_attribute.rb
index 4a3e9fc..e810a69 100644
--- a/lib/enum_machine/build_attribute.rb
+++ b/lib/enum_machine/build_attribute.rb
@@ -11,7 +11,7 @@ def self.call(enum_values:, i18n_scope:, decorator:, machine: nil)
         define_method(:machine) { machine } if machine
 
         def inspect
-          "#<EnumMachine:BuildAttribute value=#{self}>"
+          "#<EnumMachine \"#{self}\">"
         end
 
         if machine&.transitions?
diff --git a/spec/enum_machine/active_record_enum_spec.rb b/spec/enum_machine/active_record_enum_spec.rb
index aee0182..bd652dd 100644
--- a/spec/enum_machine/active_record_enum_spec.rb
+++ b/spec/enum_machine/active_record_enum_spec.rb
@@ -34,11 +34,6 @@
     expect { m.color.wrong? }.to raise_error(NoMethodError)
   end
 
-  it "pretty print inspect" do
-    m = model.new(state: "choice")
-    expect(m.state.inspect).to match(/EnumMachine:BuildAttribute.+value=choice parent=/)
-  end
-
   it "test I18n" do
     I18n.load_path = Dir["#{File.expand_path('spec/locales')}/*.yml"]
     I18n.default_locale = :ru

From 5c6aab89f5ccda4e481e5fcb08f113e8f795c8ae Mon Sep 17 00:00:00 2001
From: Michael Khabarov <michaelkhabarov@alphastate.ru>
Date: Wed, 27 Nov 2024 14:46:19 +0800
Subject: [PATCH 9/9] test: test for existing record

---
 spec/enum_machine/active_record_enum_spec.rb | 45 ++++++++++++--------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/spec/enum_machine/active_record_enum_spec.rb b/spec/enum_machine/active_record_enum_spec.rb
index bd652dd..66f8a50 100644
--- a/spec/enum_machine/active_record_enum_spec.rb
+++ b/spec/enum_machine/active_record_enum_spec.rb
@@ -78,6 +78,34 @@
     end
   end
 
+  context "when with decorator" do
+    let(:decorator_module) do
+      Module.new do
+        def am_i_choice?
+          self == "choice"
+        end
+      end
+    end
+
+    let(:model_with_decorator) do
+      decorator = decorator_module
+      Class.new(TestModel) do
+        enum_machine :state, %w[choice in_delivery], decorator: decorator
+      end
+    end
+
+    it "decorates enum value for new record" do
+      expect(model_with_decorator.new(state: "choice").state.am_i_choice?).to be(true)
+      expect(model_with_decorator.new(state: "in_delivery").state.am_i_choice?).to be(false)
+    end
+
+    it "decorates enum value for existing record" do
+      model_with_decorator.create(state: "choice")
+      m = model_with_decorator.find_by(state: "choice")
+      expect(m.state.am_i_choice?).to be(true)
+    end
+  end
+
   it "serialize model" do
     Object.const_set(:TestModelSerialize, model)
     m = TestModelSerialize.create(state: "choice", color: "wrong")
@@ -114,23 +142,6 @@
     expect(decorated_klass::COLOR::RED).to eq "red"
   end
 
-  it "keeps class of enum value" do
-    decorator =
-      Module.new do
-        def am_i_choice?
-          self == "choice"
-        end
-      end
-
-    model =
-      Class.new(TestModel) do
-        enum_machine :state, %w[choice in_delivery], decorator: decorator
-      end
-
-    expect(model.new(state: "choice").state.am_i_choice?).to be(true)
-    expect(model.new(state: "in_delivery").state.am_i_choice?).to be(false)
-  end
-
   it "returns state value by []" do
     expect(model::STATE["in_delivery"]).to eq "in_delivery"
     expect(model::STATE["in_delivery"].in_delivery?).to be(true)