app/models/hyrax/permission_template.rb in hyrax-2.0.3 vs app/models/hyrax/permission_template.rb in hyrax-2.1.0.beta1

@@ -26,14 +26,36 @@
     has_many :available_workflows, class_name: 'Sipity::Workflow', dependent: :destroy, foreign_key: :permission_template_id
 
     # In a perfect world, there would be a join table that enforced uniqueness on the ID.
     has_one :active_workflow, -> { where(active: true) }, class_name: 'Sipity::Workflow', foreign_key: :permission_template_id
 
+    # A bit of an analogue for a `belongs_to :source_model` as it crosses from Fedora to the DB
+    # @return [AdminSet | Collection]
+    # @raise [ActiveFedora::ObjectNotFoundError] when neither an AdminSet or Collection is found
+    def source_model
+      admin_set
+    rescue ActiveFedora::ObjectNotFoundError
+      collection
+    end
+
     # A bit of an analogue for a `belongs_to :admin_set` as it crosses from Fedora to the DB
     # @return [AdminSet]
     # @raise [ActiveFedora::ObjectNotFoundError] when the we cannot find the AdminSet
     def admin_set
-      AdminSet.find(admin_set_id)
+      return AdminSet.find(source_id) if AdminSet.exists?(source_id)
+      raise ActiveFedora::ObjectNotFoundError
+    rescue ActiveFedora::ActiveFedoraError # TODO: remove the rescue when active_fedora issue #1276 is fixed
+      raise ActiveFedora::ObjectNotFoundError
+    end
+
+    # A bit of an analogue for a `belongs_to :collection` as it crosses from Fedora to the DB
+    # @return [Collection]
+    # @raise [ActiveFedora::ObjectNotFoundError] when the we cannot find the Collection
+    def collection
+      return Collection.find(source_id) if Collection.exists?(source_id)
+      raise ActiveFedora::ObjectNotFoundError
+    rescue ActiveFedora::ActiveFedoraError # TODO: remove the rescue when active_fedora issue #1276 is fixed
+      raise ActiveFedora::ObjectNotFoundError
     end
 
     # Valid Release Period values
     RELEASE_TEXT_VALUE_FIXED = 'fixed'.freeze
     RELEASE_TEXT_VALUE_NO_DELAY = 'now'.freeze