6 files changed, 24 insertions, 20 deletions
diff --git a/app/controllers/api_keys_controller.rb b/app/controllers/api_keys_controller.rb
index 9706c5961..a03a67481 100644
--- a/app/controllers/api_keys_controller.rb
+++ b/app/controllers/api_keys_controller.rb
@@ -31,11 +31,4 @@ class ApiKeysController < ChouetteController
   def api_key_params
     params.require(:api_key).permit(:name, :referential_id)
   end
-
-  def decorate_api_keys(api_keys)
-    ModelDecorator.decorate(
-      api_keys,
-      with: ApiKeyDecorator,
-    )
-  end
 end
diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index 80d194096..8bd3da2f9 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -52,6 +52,19 @@ class ApplicationController < ActionController::Base
   end
   helper_method :current_functional_scope
 
+  def collection_name
+    self.class.name.split("::").last.gsub('Controller', '').underscore
+  end
+
+  def decorated_collection
+    if instance_variable_defined?("@#{collection_name}")
+      instance_variable_get("@#{collection_name}")
+    else
+      nil
+    end
+  end
+  helper_method :decorated_collection
+
   def begin_of_association_chain
     current_organisation
   end
diff --git a/app/controllers/calendars_controller.rb b/app/controllers/calendars_controller.rb
index 4a752f2b9..de8f9d10c 100644
--- a/app/controllers/calendars_controller.rb
+++ b/app/controllers/calendars_controller.rb
@@ -7,14 +7,11 @@ class CalendarsController < ChouetteController
 
   def index
     index! do
-      @calendars = ModelDecorator.decorate(@calendars, with: CalendarDecorator)
+      @calendars = CalendarDecorator.decorate(@calendars)
     end
   end
 
   def show
-    show! do
-      @calendar = @calendar.decorate
-    end
   end
 
   private
@@ -34,7 +31,7 @@ class CalendarsController < ChouetteController
 
   protected
   def resource
-    @calendar = Calendar.where('organisation_id = ? OR shared = true', current_organisation.id).find_by_id(params[:id])
+    @calendar = Calendar.where('organisation_id = ? OR shared = true', current_organisation.id).find_by_id(params[:id]).decorate
   end
 
   def build_resource
diff --git a/app/controllers/companies_controller.rb b/app/controllers/companies_controller.rb
index f84252920..4afd12be1 100644
--- a/app/controllers/companies_controller.rb
+++ b/app/controllers/companies_controller.rb
@@ -47,6 +47,9 @@ class CompaniesController < ChouetteController
     end
   end
 
+  def resource
+    super.decorate(context: { referential: line_referential })
+  end
 
   def resource_url(company = nil)
     line_referential_company_path(line_referential, company || resource)
@@ -79,9 +82,8 @@ class CompaniesController < ChouetteController
   end
 
   def decorate_companies(companies)
-    ModelDecorator.decorate(
+    CompanyDecorator.decorate(
       companies,
-      with: CompanyDecorator,
       context: {
         referential: line_referential
       }
diff --git a/app/controllers/lines_controller.rb b/app/controllers/lines_controller.rb
index 10ecc7e9d..22578187f 100644
--- a/app/controllers/lines_controller.rb
+++ b/app/controllers/lines_controller.rb
@@ -70,10 +70,6 @@ class LinesController < ChouetteController
     end
   end
 
-  def decorated_collection
-    @lines
-  end
-
   helper_method :decorated_collection
 
   protected
diff --git a/app/controllers/referential_companies_controller.rb b/app/controllers/referential_companies_controller.rb
index 7e65a72cf..806a70c8f 100644
--- a/app/controllers/referential_companies_controller.rb
+++ b/app/controllers/referential_companies_controller.rb
@@ -69,10 +69,13 @@ class ReferentialCompaniesController < ChouetteController
     %w[asc desc].include?(params[:direction]) ?  params[:direction] : 'asc'
   end
 
+  def collection_name
+    "companies"
+  end
+
   def decorate_companies(companies)
-    ModelDecorator.decorate(
+    CompanyDecorator.decorate(
       companies,
-      with: CompanyDecorator,
       context: {
         referential: referential
       }