Tree!

tree.rb

require 'pry'

Record = Struct.new(:id, :level_1, :level_2)

class Tree
  class Node
    def initialize(object, records, behavior)
      @object = object
      @records = records
      @behavior = behavior
      load_behavior
    end

    def nodes
      raise NotImplementedError
    end

    private

    def load_behavior
      name = self.class.name.gsub('Tree::', '')
      return unless @behavior.const_defined?(name)

      extend(@behavior.const_get(name))
    end
  end

  class Level1Node < Node
    def nodes
      @records.group_by(&:level_2).map do |level_2, nested|
        Level2Node.new(level_2, nested, @behavior)
      end
    end
  end

  Level2Node = Class.new(Node)

  # class methods

  def initialize(records:, behavior:)
    @records = records
    @behavior = behavior
  end

  def nodes
    @records.group_by(&:level_1).map do |level_1, nested|
      Level1Node.new(level_1, nested, @behavior)
    end
  end
end

module BarBehavior
  module Level1Node
    def custom_method
      ['This is: ', @object].join
    end
  end

  module Level2Node
    def nodes
      @records.map(&:id)
    end
  end
end


##### TEST

records = [
  Record.new(1, :foo, :a),
  Record.new(2, :foo, :a),
  Record.new(3, :bar, :b)
]

tree = Tree.new(records: records, behavior: BarBehavior)

level1 = tree.nodes.first
raise 'Does not work' unless level1.custom_method == 'This is: foo'
raise 'Does not work' unless level1.nodes.count == 1

level2 = level1.nodes.first

raise 'Does not work' unless level2.nodes == [1, 2]

puts 'WORK!'