lib/cistern/service.rb in cistern-0.2.3 vs lib/cistern/service.rb in cistern-0.3.0

@@ -26,14 +26,17 @@
 
           def service
             #{klass.name}
           end
         end
+
         def self.service
           #{klass.name}
         end
       EOS
+
+      klass.send(:const_set, :Timeout, Class.new(Cistern::Error))
     end
 
     def model_path(model_path)
       @model_path = model_path
     end
@@ -80,12 +83,12 @@
 
     def request(request_name)
       requests << request_name
     end
 
-    def collection(collection_name)
-      collections << collection_name
+    def collection(collection_name, options={})
+      collections << [collection_name, options]
     end
 
     def validate_options(options={})
       required_options = Cistern::Hash.slice(options, *required_arguments)
       missing_required_options = required_arguments - required_options.keys
@@ -121,12 +124,12 @@
                 Cistern::Mock.not_implemented
               end
             EOS
           end
         end
-        collections.each do |collection|
-          require File.join(@model_path, collection.to_s)
+        collections.each do |collection, options|
+          require File.join(@model_path, collection.to_s) unless options[:require] == false
           class_name = collection.to_s.split("_").map(&:capitalize).join
           self.const_get(:Collections).module_eval <<-EOS, __FILE__, __LINE__
             def #{collection}(attributes={})
               #{service}::#{class_name}.new({connection: self}.merge(attributes))
             end
@@ -138,16 +141,15 @@
 
     def new(options={})
       validate_options(options)
       setup_requirements
 
-      if self.mocking?
-        self.const_get(:Mock).send(:include, self.const_get(:Collections))
-        self.const_get(:Mock).new(options)
-      else
-        self.const_get(:Real).send(:include, self.const_get(:Collections))
-        self.const_get(:Real).new(options)
-      end
+      klass = self.const_get(self.mocking? ? :Mock : :Real)
+
+      klass.send(:include, service::Collections)
+      klass.send(:extend, Cistern::WaitFor)
+      klass.timeout_error = service::Timeout
+      klass.new(options)
     end
 
     def reset!
       self.const_get(:Mock).reset!
     end