lib/has_states/stateable.rb in stateful_models-0.0.1 vs lib/has_states/stateable.rb in stateful_models-0.0.2

@@ -3,17 +3,18 @@
 module HasStates
   module Stateable
     extend ActiveSupport::Concern
 
     included do
-      has_many :states, class_name: 'HasStates::State',
-                        as: :stateable,
-                        dependent: :destroy
+      has_many :states, class_name: 'HasStates::Base',
+                       as: :stateable,
+                       dependent: :destroy
     end
 
     # Instance methods for managing states
-    def add_state(type, status: 'pending', metadata: {})
+    def add_state(type, status: 'pending', metadata: {}, state_class: HasStates::State)
       states.create!(
+        type: state_class.name,
         state_type: type,
         status: status,
         metadata: metadata
       )
     end