1 files changed, 6 insertions, 5 deletions
diff --git a/lib/gitlab/graphql/representation/tree_entry.rb b/lib/gitlab/graphql/representation/tree_entry.rb
index 7ea83db5876..4467c6634ea 100644
--- a/lib/gitlab/graphql/representation/tree_entry.rb
+++ b/lib/gitlab/graphql/representation/tree_entry.rb
@@ -5,23 +5,24 @@ module Gitlab
     module Representation
       class TreeEntry < SimpleDelegator
         class << self
-          def decorate(entries, repository)
+          def decorate(entries, tree)
             return if entries.nil?
 
             entries.map do |entry|
               if entry.is_a?(TreeEntry)
                 entry
               else
-                self.new(entry, repository)
+                self.new(entry, tree)
               end
             end
           end
         end
 
-        attr_accessor :repository
+        attr_accessor :tree
+        delegate :repository, to: :tree
 
-        def initialize(raw_entry, repository)
-          @repository = repository
+        def initialize(raw_entry, tree)
+          @tree = tree
 
           super(raw_entry)
         end