lib/cistern/model.rb in cistern-0.11.1 vs lib/cistern/model.rb in cistern-0.11.2.pre2

@@ -1,11 +1,15 @@
-class Cistern::Model
-  extend Cistern::Attributes::ClassMethods
+module Cistern::Model
   include Cistern::Attributes::InstanceMethods
 
-  attr_accessor :collection, :connection
+  def self.included(klass)
+    klass.send(:extend, Cistern::Attributes::ClassMethods)
+    klass.send(:include, Cistern::Attributes::InstanceMethods)
+  end
 
+  attr_accessor :collection, :service
+
   def inspect
     if Cistern.formatter
       Cistern.formatter.call(self)
     else
       "#<#{self.class} #{self.identity}"
@@ -50,17 +54,17 @@
     else
       super
     end
   end
 
-  def service
-    self.connection ? self.connection.class : Cistern
+  def wait_for(timeout = self.service_class.timeout, interval = self.service_class.poll_interval, &block)
+    service_class.wait_for(timeout, interval) { reload && block.call(self) }
   end
 
-  def wait_for(timeout = self.service.timeout, interval = self.service.poll_interval, &block)
-    service.wait_for(timeout, interval) { reload && block.call(self) }
+  def wait_for!(timeout = self.service_class.timeout, interval = self.service_class.poll_interval, &block)
+    service_class.wait_for!(timeout, interval) { reload && block.call(self) }
   end
 
-  def wait_for!(timeout = self.service.timeout, interval = self.service.poll_interval, &block)
-    service.wait_for!(timeout, interval) { reload && block.call(self) }
+  def service_class
+    self.service ? self.service.class : Cistern
   end
 end