lib/machinist/mongoid.rb in machinist-mongoid-0.0.1 vs lib/machinist/mongoid.rb in machinist-mongoid-0.1.0

@@ -8,76 +8,87 @@
   puts "Mongoid is not installed (gem install mongoid)"
   exit
 end
 
 module Machinist
-  
+
   module Mongoid
-    
+
     module Machinable
       extend ActiveSupport::Concern
-      
+
       module ClassMethods
         include Machinist::Machinable
         def blueprint_class
           Machinist::Mongoid::Blueprint
         end
       end
     end
-    
+
     class Blueprint < Machinist::Blueprint
-      
+
       def make!(attributes = {})
         object = make(attributes)
         object.save!
         object.reload
       end
-      
+
       def lathe_class #:nodoc:
         Machinist::Mongoid::Lathe
       end
-      
+
       def outside_transaction
         yield
       end
-      
+
       def box(object)
         object.id
       end
 
       # Unbox an object from the warehouse.
       def unbox(id)
         @klass.find(id)
       end
     end
-    
+
     class Lathe < Machinist::Lathe
       def make_one_value(attribute, args) #:nodoc:
         if block_given?
           raise_argument_error(attribute) unless args.empty?
           yield
         else
           make_association(attribute, args)
         end
       end
-      
+
       def make_association(attribute, args) #:nodoc:
-        association = @klass.associations[attribute.to_s]
-        if association
-          association.klass.make(*args)
+        relation = relations_for(@klass)[attribute.to_s]
+        if relation
+          relation.klass.make(*args)
         else
           raise_argument_error(attribute)
         end
       end
-      
+
       def assign_attribute(key, value)
         @assigned_attributes[key.to_sym] = value
         if @object.respond_to?("#{key}=")
           @object.send("#{key}=", value)
         else
           @object[key] = value
         end
       end
+
+    private
+
+    def relations_for(obj)
+      if ::Mongoid::VERSION =~ /^4\./
+        obj.relations
+      else
+        obj.associations
+      end
+    end
+
     end
   end
 end
 
 module Mongoid #:nodoc: