lib/wcc/contentful/middleware/store.rb in wcc-contentful-1.3.2 vs lib/wcc/contentful/middleware/store.rb in wcc-contentful-1.4.0.rc1

@@ -31,61 +31,63 @@
     end
   end
 
   def find(id, **options)
     found = store.find(id, **options)
-    return transform(found) if found && (!has_select? || select?(found))
+    return transform(found, options) if found && (!has_select? || select?(found, options))
   end
 
   def find_by(options: nil, **args)
+    options ||= {}
     result = store.find_by(**args.merge(options: options))
-    return unless result && (!has_select? || select?(result))
+    return unless result && (!has_select? || select?(result, options))
 
-    result = resolve_includes(result, options[:include]) if options && options[:include]
-    transform(result)
+    result = resolve_includes(result, options[:include], options) if options && options[:include]
+    transform(result, options)
   end
 
   def find_all(options: nil, **args)
+    options ||= {}
     DelegatingQuery.new(
       store.find_all(**args.merge(options: options)),
       middleware: self,
       options: options
     )
   end
 
-  def resolve_includes(entry, depth)
+  def resolve_includes(entry, depth, options)
     return entry unless entry && depth && depth > 0
 
     # We only care about entries (see #resolved_link?)
-    WCC::Contentful::LinkVisitor.new(entry, :Entry, depth: depth).map! do |val|
-      resolve_link(val)
+    WCC::Contentful::LinkVisitor.new(entry, :Entry, :Asset, depth: depth).map! do |val|
+      resolve_link(val, options)
     end
   end
 
-  def resolve_link(val)
+  def resolve_link(val, options)
     return val unless resolved_link?(val)
 
-    if !has_select? || select?(val)
-      transform(val)
+    if !has_select? || select?(val, options)
+      transform(val, options)
     else
       # Pretend it's an unresolved link -
       # matches the behavior of a store when the link cannot be retrieved
       WCC::Contentful::Link.new(val.dig('sys', 'id'), val.dig('sys', 'type')).to_h
     end
   end
 
   def resolved_link?(value)
-    value.is_a?(Hash) && value.dig('sys', 'type') == 'Entry'
+    value.is_a?(Hash) && %w[Entry Asset].include?(value.dig('sys', 'type'))
   end
 
   def has_select? # rubocop:disable Naming/PredicateName
     respond_to?(:select?)
   end
 
   # The default version of `#transform` just returns the entry.
   # Override this with your own implementation.
-  def transform(entry)
+  def transform(entry, _options)
     entry
   end
 
   class DelegatingQuery
     include WCC::Contentful::Store::Query::Interface
@@ -108,15 +110,17 @@
 
     attr_reader :wrapped_query, :middleware, :options
 
     def to_enum
       result = wrapped_query.to_enum
-      result = result.select { |x| middleware.select?(x) } if middleware.has_select?
+      result = result.select { |x| middleware.select?(x, options) } if middleware.has_select?
 
-      result = result.map { |x| middleware.resolve_includes(x, options[:include]) } if options && options[:include]
+      if options && options[:include]
+        result = result.map { |x| middleware.resolve_includes(x, options[:include], options) }
+      end
 
-      result.map { |x| middleware.transform(x) }
+      result.map { |x| middleware.transform(x, options) }
     end
 
     def apply(filter, context = nil)
       self.class.new(
         wrapped_query.apply(filter, context),
@@ -148,10 +152,10 @@
     end
 
     def initialize(wrapped_query, middleware:, options: nil, **extra)
       @wrapped_query = wrapped_query
       @middleware = middleware
-      @options = options
+      @options = options || {}
       @extra = extra
     end
   end
 end