xcknife 0.1.0 → 0.5.0

data/lib/xcknife/exceptions.rb

@@ -0,0 +1,6 @@
+module XCKnife
+  # Base error class for xcknife
+  XCKnifeError = Class.new(StandardError)
+
+  StreamParsingError = Class.new(XCKnifeError)
+end

data/lib/xcknife/json_stream_parser_helper.rb

@@ -0,0 +1,24 @@
+module XCKnife
+  module JsonStreamParserHelper
+    extend self
+
+    # Iterates over events, calling block once for each test_target/test event on a events (a parsed json_stream iterable)
+    def each_test_event(events, &block)
+      current_target = nil
+      events.each do |result|
+        current_target = result.targetName if result.event == "begin-ocunit"
+        if result.test and result.event == "end-test"
+          raise XCKnife::StreamParsingError, "No test target defined" if current_target.nil?
+          block.call(current_target, normalize_result(result))
+        end
+      end
+    end
+
+    def normalize_result(result)
+      if result.totalDuration.is_a?(String)
+        result.totalDuration = result.totalDuration.to_f
+      end
+      result
+    end
+  end
+end

data/lib/xcknife/runner.rb

@@ -0,0 +1,138 @@
+require 'optparse'
+
+module XCKnife
+  class Runner
+    include XCToolCmdHelper
+
+    def self.invoke
+      new(ARGV).run
+    end
+
+    attr_reader :parser
+
+    def initialize(args)
+      @abbreviated_output = false
+      @partitions = []
+      @partition_names = []
+      @worker_count = nil
+      @historical_timings_file = nil
+      @current_tests_file = nil
+      @output_file_name = nil
+      parse_arguments(args)
+    end
+
+    def run
+      stream_parser = XCKnife::StreamParser.new(@worker_count, @partitions)
+      result = stream_parser.compute_shards_for_file(@historical_timings_file, @current_tests_file)
+      data = @abbreviated_output ? gen_abbreviated_output(result) : gen_full_output(result)
+      write_output(data)
+    rescue XCKnife::XCKnifeError => e
+      warn "Error: #{e}"
+      exit 1
+    end
+    private
+    def gen_abbreviated_output(result)
+      result.test_maps.map { |partition_set| xctool_only_arguments_for_a_partition_set(partition_set) }
+    end
+
+    def gen_full_output(result)
+      {
+        metadata: {
+          worker_count: @worker_count,
+          partition_set_count: result.test_maps.size,
+          total_time_in_ms: result.total_test_time,
+        }.merge(result.stats.to_h),
+        partition_set_data: partition_sets_data(result)
+      }
+    end
+
+    def partition_sets_data(result)
+      shard_number = 0
+      result.test_maps.each_with_index.map do |partition_set, partition_set_i|
+        partition_data = partition_set.each_with_index.map do |partition, partition_j|
+          shard_number += 1
+          partition_data(result, shard_number, partition, partition_set_i, partition_j)
+        end
+
+        {
+          partition_set: @partition_names[partition_set_i],
+          size: partition_set.size,
+          imbalance_ratio: result.test_time_imbalances.partition_set[partition_set_i],
+          partitions: partition_data
+        }
+      end
+    end
+
+    def partition_data(result, shard_number, partition, partition_set_i, partition_j)
+      {
+        shard_number: shard_number,
+        cli_arguments: xctool_only_arguments(partition),
+        partition_imbalance_ratio: result.test_time_imbalances.partitions[partition_set_i][partition_j]
+      }
+    end
+
+    def write_output(data)
+      json = JSON.pretty_generate(data)
+      return puts json if @output_file_name.nil?
+      File.open(@output_file_name, "w") { |f| f.puts(json) }
+      puts "Wrote file to: #{@output_file_name}"
+    end
+
+    def parse_arguments(args)
+      positional_arguments = parse_options(args)
+      if positional_arguments.size < required_arguments.size
+        warn_and_exit("You must specify *all* required arguments: #{required_arguments.join(", ")}")
+      end
+      if @partitions.empty?
+        warn_and_exit("At least one target partition set must be provided with -p flag")
+      end
+      worker_count, @historical_timings_file, @current_tests_file = positional_arguments
+      @worker_count = Integer(worker_count)
+    end
+
+    def parse_options(args)
+      build_parser
+      begin
+        parser.parse(args)
+      rescue OptionParser::ParseError => error
+        warn_and_exit(error)
+      end
+    end
+
+    def build_parser
+      @parser = OptionParser.new do |opts|
+        opts.banner += " #{arguments_banner}"
+        opts.on("-p", "--partition TARGETS",
+          "Comma separated list of targets. Can be used multiple times.") do |v|
+          @partition_names << v
+          @partitions << v.split(",")
+        end
+        opts.on("-o", "--output FILENAME", "Output file. Defaults to STDOUT") { |v| @output_file_name = v }
+        opts.on("-a", "--abbrev", "Results are abbreviated") { |v| @abbreviated_output = v }
+
+        opts.on_tail("-h", "--help", "Show this message") do
+          puts opts
+          exit
+        end
+      end
+    end
+
+    def required_arguments
+      %w[worker-count historical-timings-json-stream-file]
+    end
+
+    def optional_arguments
+      %w[current-tests-json-stream-file]
+    end
+
+    def arguments_banner
+      optional_args = optional_arguments.map { |a| "[#{a}]" }
+      (required_arguments + optional_args).join(" ")
+    end
+
+    def warn_and_exit(msg)
+      warn "#{msg.to_s.capitalize} \n\n#{parser}"
+      exit 1
+    end
+  end
+end

data/lib/xcknife/stream_parser.rb

@@ -0,0 +1,222 @@
+require 'xcknife/json_stream_parser_helper'
+require 'json'
+require 'set'
+require 'ostruct'
+require 'forwardable'
+
+module XCKnife
+  class StreamParser
+    include JsonStreamParserHelper
+
+    attr_reader :number_of_shards, :test_partitions, :stats, :relevant_partitions
+
+    def initialize(number_of_shards, test_partitions)
+      @number_of_shards = number_of_shards
+      @test_partitions = test_partitions.map(&:to_set)
+      @relevant_partitions = test_partitions.flatten.to_set
+      @stats = ResultStats.new
+      ResultStats.members.each { |k| @stats[k] = 0}
+    end
+
+    PartitionWithMachines = Struct.new :test_time_map, :number_of_shards, :partition_time, :max_shard_count
+    MachineAssignment = Struct.new :test_time_map, :total_time
+    ResultStats = Struct.new :historical_total_tests, :current_total_tests, :class_extrapolations, :target_extrapolations
+
+    class PartitionResult
+      TimeImbalances = Struct.new :partition_set, :partitions
+      attr_reader :stats, :test_maps, :test_times, :total_test_time, :test_time_imbalances
+      extend Forwardable
+      delegate ResultStats.members => :@stats
+
+      def initialize(stats, partition_sets)
+        @stats = stats
+        @partition_sets = partition_sets
+        @test_maps = partition_sets_map(&:test_time_map)
+        @test_times = partition_sets_map(&:total_time)
+        @total_test_time = test_times.flatten.inject(:+)
+        @test_time_imbalances = compute_test_time_imbalances
+      end
+
+      private
+      # Yields the imbalances ratios of the partition sets, and the internal imbalance ratio of the respective partitions
+      def compute_test_time_imbalances
+        times = test_times
+        average_partition_size = times.map { |l| l.inject(:+).to_f / l.size}
+        ideal_partition_set_avg = average_partition_size.inject(:+) / @partition_sets.size
+        partition_set_imbalance = average_partition_size.map { |avg| avg / ideal_partition_set_avg }
+
+        internal_partition_imbalance = times.map do |partition_times|
+          internal_total =  partition_times.inject(:+)
+          partition_times.map do |partition_time|
+            (partition_time * partition_times.size).to_f / internal_total
+          end
+        end
+        TimeImbalances.new partition_set_imbalance, internal_partition_imbalance
+      end
+
+      def partition_sets_map(&block)
+        @partition_sets.map { |assignemnt_list| assignemnt_list.map(&block) }
+      end
+    end
+
+    # Parses the output of a xctool json-stream reporter and compute the shards based of that
+    # see: https://github.com/facebook/xctool#included-reporters
+    #
+    # @param historical_filename: String  the path of the, usually historical, test time performance.
+    # @param current_test_filename: [String, nil] = the path of the current test names and targets,
+    def compute_shards_for_file(historical_filename, current_test_filename = nil)
+      compute_shards_for_events(parse_json_stream_file(historical_filename), parse_json_stream_file(current_test_filename))
+    end
+
+    def compute_shards_for_events(historical_events, current_events = nil)
+      compute_shards_for_partitions(test_time_for_partitions(historical_events, current_events))
+    end
+
+    def compute_shards_for_partitions(test_time_for_partitions)
+      PartitionResult.new(@stats, split_machines_proportionally(test_time_for_partitions).map do |partition|
+        compute_single_shards(partition.number_of_shards, partition.test_time_map)
+      end)
+    end
+
+    def test_time_for_partitions(historical_events, current_events = nil)
+      analyzer =  EventsAnalyzer.for(current_events, relevant_partitions)
+      @stats[:current_total_tests] = analyzer.total_tests
+      times_for_target_class = Hash.new { |h, current_target| h[current_target] = Hash.new(0) }
+      each_test_event(historical_events) do |target_name, result|
+        next unless relevant_partitions.include?(target_name)
+        inc_stat :historical_total_tests
+        next unless analyzer.has_test_class?(target_name, result.className)
+        times_for_target_class[target_name][result.className] += (result.totalDuration * 1000).ceil
+      end
+
+      extrapolate_times_for_current_events(analyzer, times_for_target_class) if current_events
+      hash_partitions(times_for_target_class)
+    end
+
+    def split_machines_proportionally(partitions)
+      total = 0
+      partitions.each do |test_time_map|
+        each_duration(test_time_map) { |duration_in_milliseconds| total += duration_in_milliseconds}
+      end
+
+      used_shards = 0
+      assignable_shards = number_of_shards - partitions.size
+      partition_with_machines_list = partitions.map do |test_time_map|
+        partition_time = 0
+        max_shard_count = test_time_map.values.map(&:size).inject(:+) || 1
+        each_duration(test_time_map) { |duration_in_milliseconds| partition_time += duration_in_milliseconds}
+        n = [1 + (assignable_shards * partition_time.to_f / total).floor, max_shard_count].min
+        used_shards += n
+        PartitionWithMachines.new(test_time_map, n, partition_time, max_shard_count)
+      end
+
+      fifo_with_machines_who_can_use_more_shards = partition_with_machines_list.select { |x| x.number_of_shards < x.max_shard_count}.sort_by(&:partition_time)
+      while (number_of_shards - used_shards) > 0
+        if fifo_with_machines_who_can_use_more_shards.empty?
+          raise XCKnife::XCKnifeError, "There are #{number_of_shards - used_shards} extra machines"
+        end
+        machine = fifo_with_machines_who_can_use_more_shards.pop
+        machine.number_of_shards += 1
+        used_shards += 1
+        if machine.number_of_shards < machine.max_shard_count
+          fifo_with_machines_who_can_use_more_shards.unshift(machine)
+        end
+      end
+      partition_with_machines_list
+    end
+
+    # Computes  a 2-aproximation to the optimal partition_time, which is an instance of the Open shop scheduling problem (which is NP-hard)
+    # see: https://en.wikipedia.org/wiki/Open-shop_scheduling
+    def compute_single_shards(number_of_shards, test_time_map)
+      raise XCKnife::XCKnifeError, "There are not enough workers provided" if number_of_shards <= 0
+      raise XCKnife::XCKnifeError, "Cannot shard an empty partition_time" if test_time_map.empty?
+      assignements = Array.new(number_of_shards) { MachineAssignment.new(Hash.new { |k, v| k[v] = [] }, 0) }
+
+      list_of_test_target_class_times = []
+      test_time_map.each do |test_target, class_times|
+        class_times.each do |class_name, duration_in_milliseconds|
+          list_of_test_target_class_times << [test_target, class_name, duration_in_milliseconds]
+        end
+      end
+
+      list_of_test_target_class_times.sort_by! { |test_target, class_name, duration_in_milliseconds| -duration_in_milliseconds }
+      list_of_test_target_class_times.each do |test_target, class_name, duration_in_milliseconds|
+        assignemnt = assignements.min_by(&:total_time)
+        assignemnt.test_time_map[test_target] << class_name
+        assignemnt.total_time += duration_in_milliseconds
+      end
+      raise XCKnife::XCKnifeError, "Too many shards" if assignements.any? { |a| a.test_time_map.empty? }
+      assignements
+    end
+
+    def parse_json_stream_file(filename)
+      return nil if filename.nil?
+      return [] unless File.exists?(filename)
+      lines = IO.readlines(filename)
+      lines.lazy.map { |line| OpenStruct.new(JSON.load(line)) }
+    end
+
+    private
+    def inc_stat(name)
+      @stats[name] += 1
+    end
+
+    def each_duration(test_time_map, &block)
+      test_time_map.each do |test_target, class_times|
+        class_times.each do |class_name, duration_in_milliseconds|
+          yield(duration_in_milliseconds)
+        end
+      end
+    end
+
+    def extrapolate_times_for_current_events(analyzer, times_for_target_class)
+      median_map = {}
+      times_for_target_class.each do |test_target, class_times|
+        median_map[test_target] = median(class_times.values)
+      end
+
+      all_times_for_all_classes = times_for_target_class.values.flat_map(&:values)
+      median_of_targets = median(all_times_for_all_classes)
+      analyzer.target_class_map.each do |test_target, class_set|
+        if times_for_target_class.has_key?(test_target)
+          class_set.each do |clazz|
+            unless times_for_target_class[test_target].has_key?(clazz)
+              inc_stat :class_extrapolations
+              times_for_target_class[test_target][clazz] = median_map[test_target]
+            end
+          end
+        else
+          inc_stat :target_extrapolations
+          class_set.each do |clazz|
+            inc_stat :class_extrapolations
+            times_for_target_class[test_target][clazz] = extrapolated_duration(median_of_targets, class_set)
+          end
+        end
+      end
+    end
+
+    DEFAULT_EXTRAPOLATED_DURATION = 1000
+    def extrapolated_duration(median_of_targets, class_set)
+      return DEFAULT_EXTRAPOLATED_DURATION if median_of_targets.nil?
+      median_of_targets / class_set.size
+    end
+
+    def median(array)
+      array.sort[array.size / 2]
+    end
+
+    def hash_partitions(times)
+      ret = Array.new(test_partitions.size) { {} }
+      times.each do |test_target, times_map|
+        test_partitions.each_with_index do |partition, i|
+          ret[i][test_target] = times_map if partition.include?(test_target)
+        end
+      end
+      ret.each_with_index do |partition, index|
+        if partition.empty?
+          raise XCKnife::XCKnifeError, "The following partition has no tests: #{test_partitions[index].to_a.inspect}"
+        end
+      end
+    end
+  end
+end

data/lib/xcknife/xctool_cmd_helper.rb

@@ -0,0 +1,14 @@
+require 'pp'
+module XCKnife
+  module XCToolCmdHelper
+    def xctool_only_arguments(single_partition)
+      single_partition.flat_map do |test_target, classes|
+        ['-only', "#{test_target}:#{classes.sort.join(',')}"]
+      end
+    end
+
+    def xctool_only_arguments_for_a_partition_set(partition_set)
+      partition_set.map { |partition| xctool_only_arguments(partition) }
+    end
+  end
+end

data/spec/spec_helper.rb

@@ -0,0 +1,4 @@
+$LOAD_PATH.unshift File.expand_path('../../lib', __FILE__)
+require 'xcknife'
+require 'set'
+require 'pp'

data/spec/xcknife_spec.rb

@@ -0,0 +1,344 @@
+require 'spec_helper'
+
+describe XCKnife::StreamParser do
+  context 'test_time_for_partitions' do
+    subject { XCKnife::StreamParser.new(2, [["TestTarget1"], ["TestTarget2"]]) }
+
+    it 'decide how many shards each partition set needs' do
+      stream = [xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1"),
+        xctool_target_event("TestTarget2"),
+        xctool_test_event("ClassTest2", "test1")
+      ]
+      result = subject.test_time_for_partitions(stream)
+      expect(result).to eq([{ "TestTarget1" => { "ClassTest1" => 1000 } },
+        { "TestTarget2" => { "ClassTest2" => 1000 } }])
+    end
+
+    it 'aggretates the times at the class level' do
+      stream_parser = XCKnife::StreamParser.new(2, [["TestTarget1"]])
+      stream = [xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1", 1.0),
+        xctool_test_event("ClassTest1", "test2", 2.0)
+      ]
+      result = stream_parser.test_time_for_partitions(stream)
+      expect(result).to eq([{ "TestTarget1" => { "ClassTest1" => 3000 } }])
+    end
+
+    it 'works with multiple partitions' do
+      stream_parser = XCKnife::StreamParser.new(2, [["TestTarget1"], ["TestTarget2"], ["TestTarget3"]])
+
+      stream = [xctool_target_event("TestTarget1"),
+        xctool_test_event("Class1", "test1"),
+        xctool_target_event("TestTarget2"),
+        xctool_test_event("Class2", "test1"),
+        xctool_target_event("TestTarget3"),
+        xctool_test_event("Class3", "test1"),
+      ]
+      result = stream_parser.test_time_for_partitions(stream)
+      expect(result).to eq([{ "TestTarget1" => { "Class1" => 1000 } },
+        { "TestTarget2" => { "Class2" => 1000 } },
+        { "TestTarget3" => { "Class3" => 1000 } }])
+    end
+
+    it 'allows the same target to be listed on multiple partitions' do
+      stream_parser = XCKnife::StreamParser.new(2, [["TestTarget1"], ["TestTarget2", "TestTarget1"]])
+      stream = [xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1"),
+        xctool_target_event("TestTarget2"),
+        xctool_test_event("ClassTest2", "test1"),
+      ]
+      result = stream_parser.test_time_for_partitions(stream)
+      expect(result).to eq([{ "TestTarget1" => { "ClassTest1" => 1000 } },
+        { "TestTarget2" => { "ClassTest2" => 1000 },
+          "TestTarget1" => { "ClassTest1" => 1000 }
+        }])
+    end
+
+    it 'raises error when an empty partition is specified' do
+      stream_parser = XCKnife::StreamParser.new(1, [["TestTarget1"]])
+      expect { stream_parser.test_time_for_partitions([]) }.to raise_error(XCKnife::XCKnifeError, 'The following partition has no tests: ["TestTarget1"]')
+    end
+  end
+
+  context 'provided historical events' do
+    subject { XCKnife::StreamParser.new(2, [["TestTarget1", "TestTarget2", "TestTarget3", "NewTestTarget1"]]) }
+
+    it 'ignores test targets not present on current events' do
+      historical_events = [xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1"),
+        xctool_test_event("ClassTest1", "test2"),
+        xctool_target_event("TestTarget2"),
+        xctool_test_event("ClassTest2", "test1"),
+        xctool_test_event("ClassTest2", "test2")
+      ]
+      current_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1")
+      ]
+      result = subject.test_time_for_partitions(historical_events, current_events)
+      expect(result).to eq([{ "TestTarget1" => { "ClassTest1" => 2000 } }])
+      expect(subject.stats.to_h).to eq({historical_total_tests: 4, current_total_tests: 1, class_extrapolations: 0, target_extrapolations: 0})
+    end
+
+
+    it 'ignores test classes not present on current events' do
+      historical_events = [xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1"),
+        xctool_test_event("ClassTest1", "test2"),
+        xctool_test_event("ClassTest2", "test1"),
+        xctool_test_event("ClassTest2", "test2")
+      ]
+      current_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1")
+      ]
+      result = subject.test_time_for_partitions(historical_events, current_events)
+      expect(result).to eq([{ "TestTarget1" => { "ClassTest1" => 2000 } }])
+      expect(subject.stats.to_h).to eq({historical_total_tests: 4, current_total_tests: 1, class_extrapolations: 0, target_extrapolations: 0})
+    end
+
+    it 'extrapolates for new test targets' do
+      historical_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1")
+      ]
+      current_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1"),
+        xctool_target_event("NewTestTargetButNotRelevant"),
+        xctool_test_event("ClassTest10", "test1")
+      ]
+      result = subject.test_time_for_partitions(historical_events, current_events)
+      expect(result.to_set).to eq([{
+        "TestTarget1" => { "ClassTest1" => 1000 }
+      }
+      ].to_set)
+      expect(subject.stats.to_h).to eq({historical_total_tests: 1, current_total_tests: 1, class_extrapolations: 0, target_extrapolations: 0})
+    end
+
+    it 'extrapolates for new test classes' do
+      historical_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1", 1.0),
+        xctool_test_event("ClassTest2", "test2", 5.0),
+        xctool_test_event("ClassTest3", "test3", 10000.0)
+      ]
+      current_events = historical_events + [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest2", "test2"),
+        xctool_test_event("ClassTestNew", "test1")
+      ]
+      result = subject.test_time_for_partitions(historical_events, current_events)
+      median = 5000
+      expect(result).to eq([{
+        "TestTarget1" =>
+          {
+            "ClassTest1" => 1000,
+            "ClassTest2" => 5000,
+            "ClassTest3" => 10000000,
+            "ClassTestNew" => median
+          },
+      }])
+      expect(subject.stats.to_h).to eq({historical_total_tests: 3, current_total_tests: 5, class_extrapolations: 1, target_extrapolations: 0})
+    end
+
+    it "ignores test classes that don't belong to relevant targets" do
+      historical_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1", 1.0),
+        xctool_test_event("ClassTest2", "test2", 5.0),
+        xctool_test_event("ClassTest3", "test3", 10000.0)
+      ]
+      current_events = historical_events + [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest2", "test2"),
+        xctool_test_event("ClassTestNew", "test1")
+      ]
+      result = subject.test_time_for_partitions(historical_events, current_events)
+      median = 5000
+      expect(result).to eq([{
+        "TestTarget1" =>
+          {
+            "ClassTest1" => 1000,
+            "ClassTest2" => 5000,
+            "ClassTest3" => 10000000,
+            "ClassTestNew" => median
+          },
+      }])
+      expect(subject.stats.to_h).to eq({historical_total_tests: 3, current_total_tests: 5, class_extrapolations: 1, target_extrapolations: 0})
+    end
+  end
+
+  context 'provided an empty set of applicable historical events' do
+    subject { XCKnife::StreamParser.new(2, [["TestTarget1", "TestTarget2", "TestTarget3", "NewTestTarget1"]]) }
+
+    let(:empty_historical_events) { [] }
+    let(:default_extrapolated_duration) { XCKnife::StreamParser::DEFAULT_EXTRAPOLATED_DURATION }
+
+    it 'extrapolates the test target duration and classes get extrapolated' do
+      current_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1")
+      ]
+      result = subject.test_time_for_partitions(empty_historical_events, current_events)
+      expect(result).to eq([{ "TestTarget1" => { "ClassTest1" => default_extrapolated_duration } }])
+    end
+
+    it 'extrapolates the test target to different classes' do
+      effectively_empty_historical_events = [
+        xctool_target_event("TestTarget2"),
+        xctool_test_event("IgnoredClass", "ignoredTest"),
+      ]
+      current_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest1", "test1"),
+        xctool_test_event("ClassTest2", "test2")
+      ]
+      result = subject.test_time_for_partitions(effectively_empty_historical_events, current_events)
+      duration = default_extrapolated_duration
+      expect(result).to eq([{ "TestTarget1" => { "ClassTest1" => duration, "ClassTest2" => duration } }])
+    end
+
+    it "can handle multiple test targets and test classes" do
+      current_events = [
+        xctool_target_event("TestTarget1"),
+        xctool_test_event("ClassTest11", "test1"),
+        xctool_test_event("ClassTest12", "test1"),
+        xctool_test_event("ClassTest13", "test1"),
+        xctool_target_event("TestTarget2"),
+        xctool_test_event("ClassTest21", "test1"),
+        xctool_test_event("ClassTest22", "test1"),
+        xctool_test_event("ClassTest23", "test1"),
+        xctool_target_event("TestTarget3"),
+        xctool_test_event("ClassTest31", "test1"),
+        xctool_test_event("ClassTest32", "test1"),
+        xctool_test_event("ClassTest33", "test1"),
+      ]
+      result = subject.test_time_for_partitions(empty_historical_events, current_events)
+      duration = default_extrapolated_duration
+      expect(result).to eq(
+        [
+          {
+            "TestTarget1" => { "ClassTest11" => duration, "ClassTest12" => duration, "ClassTest13" => duration },
+            "TestTarget2" => { "ClassTest21" => duration, "ClassTest22" => duration, "ClassTest23" => duration },
+            "TestTarget3" => { "ClassTest31" => duration, "ClassTest32" => duration, "ClassTest33" => duration } }]
+      )
+    end
+  end
+
+  it "can split_machines_proportionally" do
+    stream_parser = XCKnife::StreamParser.new(5, [["TargetOnPartition1"], ["TargetOnPartition2"]])
+    result = stream_parser.split_machines_proportionally([
+      { "TargetOnPartition1" => { "TestClass1" => 500, "TestClass2" => 500 } },
+      { "TargetOnPartition2" => { "TestClass3" => 1000, "TestClass4" => 1000, "TestClass5" => 1000, "TestClass6" => 1000 } }])
+    expect(result.map(&:number_of_shards)).to eq([1, 4])
+  end
+
+  it "can split_machines_proportionally even when in the presence of large imbalances" do
+    stream_parser = XCKnife::StreamParser.new(5, [["TargetOnPartition1"], ["TargetOnPartition2"], ["TargetOnPartition3"]])
+    result = stream_parser.split_machines_proportionally([{ "TargetOnPartition1" => { "TestClass1" => 1 } },
+      { "TargetOnPartition2" => { "TestClass2" => 1} },
+      { "TargetOnPartition3" => { "TestClass3" => 1000, "TestClass4" => 1000, "TestClass5" => 1000} }])
+    expect(result.map(&:number_of_shards)).to eq([1, 1, 3])
+  end
+
+
+  it "should never let partition_sets have less than 1 machine alocated to them" do
+    stream_parser = XCKnife::StreamParser.new(3, [["TestTarget1"], ["TestTarget2"]])
+    result = stream_parser.split_machines_proportionally([{ "TargetOnPartition1" => { "TestClass1" => 1 } },
+      { "TargetOnPartition2" => { "TestClass2" => 2000, "TestClass3" => 2000 } }])
+    expect(result.map(&:number_of_shards)).to eq([1, 2])
+  end
+
+
+  context 'test_time_for_partitions' do
+    it "partitions the test classes accross the number of machines" do
+      stream_parser = XCKnife::StreamParser.new(2, [["Test Target"]])
+      partition = { "Test Target" => { "Class1" => 1000, "Class2" => 1000, "Class3" => 2000 } }
+      shards = stream_parser.compute_single_shards(2, partition).map(&:test_time_map)
+      expect(shards.size).to eq 2
+      first_shard, second_shard = shards.sort_by { |map| map.values.flatten.size }
+      expect(first_shard.keys).to eq(["Test Target"])
+      expect(first_shard.values).to eq([["Class3"]])
+
+      expect(second_shard.keys).to eq(["Test Target"])
+      expect(second_shard.values.map(&:to_set)).to eq([["Class1", "Class2"].to_set])
+    end
+
+    it "partitions the test, across targets" do
+      stream_parser = XCKnife::StreamParser.new(2, [["Test Target1", "Test Target2", "Test Target3"]])
+      partition = { "Test Target1" => { "Class1" => 1000 },
+        "Test Target2" => { "Class2" => 1000 },
+        "Test Target3" => { "Class3" => 2000 } }
+      shards = stream_parser.compute_single_shards(2, partition).map(&:test_time_map)
+      expect(shards.size).to eq 2
+      first_shard, second_shard = shards.sort_by { |map| map.values.flatten.size }
+      expect(first_shard.keys).to eq(["Test Target3"])
+      expect(first_shard.values).to eq([["Class3"]])
+
+      expect(second_shard.keys.to_set).to eq(["Test Target1", "Test Target2"].to_set)
+      expect(second_shard.values.to_set).to eq([["Class1"], ["Class2"]].to_set)
+    end
+
+    it "raises an error if there are too many shards" do
+      too_many_machines = 2
+      stream_parser = XCKnife::StreamParser.new(too_many_machines, [["Test Target1"]])
+      partition = { "Test Target1" => { "Class1" => 1000 } }
+      expect { stream_parser.compute_single_shards(too_many_machines, partition) }.
+        to raise_error(XCKnife::XCKnifeError, "Too many shards")
+    end
+  end
+
+  it "can compute test for all partitions" do
+    stream_parser = XCKnife::StreamParser.new(3, [["TargetOnPartition1"], ["TargetOnPartition2"]])
+    result = stream_parser.compute_shards_for_partitions([{ "TargetOnPartition1" => { "TestClass1" => 1000 } },
+      { "TargetOnPartition2" => { "TestClass2" => 4000, "TestClass3" => 4000 } }])
+    expect(result.test_maps).to eq([[{ "TargetOnPartition1" => ["TestClass1"] }],
+      [{ "TargetOnPartition2" => ["TestClass2"] },
+        { "TargetOnPartition2" => ["TestClass3"] }]])
+    expect(result.test_times).to eq [[1000], [4000, 4000]]
+    expect(result.total_test_time).to eq 9000
+    expect(result.test_time_imbalances.to_h).to eq({
+      partition_set: [0.4, 1.6],
+      partitions: [[1.0], [1.0, 1.0]]
+    })
+  end
+
+  it "can compute for only one partition set" do
+    stream_parser = XCKnife::StreamParser.new(1, [["TargetOnPartition1"]])
+    historical_events = [xctool_target_event("TargetOnPartition1"),
+      xctool_test_event("ClassTest1", "test1"),
+    ]
+    result = stream_parser.compute_shards_for_events(historical_events)
+    expect(result.test_maps).to eq([[{ "TargetOnPartition1" => ["ClassTest1"] }]])
+    expect(result.test_times).to eq [[1000]]
+    expect(result.total_test_time).to eq 1000
+    expect(result.stats.to_h).to eq({historical_total_tests: 1, current_total_tests: 0, class_extrapolations: 0, target_extrapolations: 0})
+    expect(result.test_time_imbalances.to_h).to eq({
+      partition_set: [1.0],
+      partitions: [[1.0]]
+    })
+  end
+
+  def xctool_test_event(class_name, method_name, duration = 1.0)
+    OpenStruct.new({ result: "success",
+      exceptions: [],
+      test: "-[#{class_name} #{method_name}]",
+      className: class_name,
+      event: "end-test",
+      methodName: method_name,
+      succeeded: true,
+      output: "",
+      totalDuration: duration,
+      timestamp: 0
+    })
+  end
+
+  def xctool_target_event(target_name)
+    OpenStruct.new({ result: "success",
+      event: "begin-ocunit",
+      targetName: target_name
+    })
+  end
+end