4 files changed, 93 insertions, 49 deletions

diff --git a/lib/rspec_flaky/config.rb b/lib/rspec_flaky/config.rb
new file mode 100644
index 00000000000..a17ae55910e
--- /dev/null
+++ b/lib/rspec_flaky/config.rb
@@ -0,0 +1,21 @@
+require 'json'
+
+module RspecFlaky
+  class Config
+    def self.generate_report?
+      ENV['FLAKY_RSPEC_GENERATE_REPORT'] == 'true'
+    end
+
+    def self.suite_flaky_examples_report_path
+      ENV['SUITE_FLAKY_RSPEC_REPORT_PATH'] || Rails.root.join("rspec_flaky/suite-report.json")
+    end
+
+    def self.flaky_examples_report_path
+      ENV['FLAKY_RSPEC_REPORT_PATH'] || Rails.root.join("rspec_flaky/report.json")
+    end
+
+    def self.new_flaky_examples_report_path
+      ENV['NEW_FLAKY_RSPEC_REPORT_PATH'] || Rails.root.join("rspec_flaky/new-report.json")
+    end
+  end
+end
diff --git a/lib/rspec_flaky/flaky_example.rb b/lib/rspec_flaky/flaky_example.rb
index f81fb90e870..6be24014d89 100644
--- a/lib/rspec_flaky/flaky_example.rb
+++ b/lib/rspec_flaky/flaky_example.rb
@@ -9,24 +9,21 @@ module RspecFlaky
           line: example.line,
           description: example.description,
           last_attempts_count: example.attempts,
-          flaky_reports: 1)
+          flaky_reports: 0)
       else
         super
       end
     end
 
-    def first_flaky_at
-      self[:first_flaky_at] || Time.now
-    end
-
-    def last_flaky_at
-      Time.now
-    end
+    def update_flakiness!(last_attempts_count: nil)
+      self.first_flaky_at ||= Time.now
+      self.last_flaky_at = Time.now
+      self.flaky_reports += 1
+      self.last_attempts_count = last_attempts_count if last_attempts_count
 
-    def last_flaky_job
-      return unless ENV['CI_PROJECT_URL'] && ENV['CI_JOB_ID']
-
-      "#{ENV['CI_PROJECT_URL']}/-/jobs/#{ENV['CI_JOB_ID']}"
+      if ENV['CI_PROJECT_URL'] && ENV['CI_JOB_ID']
+        self.last_flaky_job = "#{ENV['CI_PROJECT_URL']}/-/jobs/#{ENV['CI_JOB_ID']}"
+      end
     end
 
     def to_h
diff --git a/lib/rspec_flaky/flaky_examples_collection.rb b/lib/rspec_flaky/flaky_examples_collection.rb
new file mode 100644
index 00000000000..973c95b0212
--- /dev/null
+++ b/lib/rspec_flaky/flaky_examples_collection.rb
@@ -0,0 +1,37 @@
+require 'json'
+
+module RspecFlaky
+  class FlakyExamplesCollection < SimpleDelegator
+    def self.from_json(json)
+      new(JSON.parse(json))
+    end
+
+    def initialize(collection = {})
+      unless collection.is_a?(Hash)
+        raise ArgumentError, "`collection` must be a Hash, #{collection.class} given!"
+      end
+
+      collection_of_flaky_examples =
+        collection.map do |uid, example|
+          [
+            uid,
+            example.is_a?(RspecFlaky::FlakyExample) ? example : RspecFlaky::FlakyExample.new(example)
+          ]
+        end
+
+      super(Hash[collection_of_flaky_examples])
+    end
+
+    def to_report
+      Hash[map { |uid, example| [uid, example.to_h] }].deep_symbolize_keys
+    end
+
+    def -(other)
+      unless other.respond_to?(:key)
+        raise ArgumentError, "`other` must respond to `#key?`, #{other.class} does not!"
+      end
+
+      self.class.new(reject { |uid, _| other.key?(uid) })
+    end
+  end
+end
diff --git a/lib/rspec_flaky/listener.rb b/lib/rspec_flaky/listener.rb
index ec2fbd9e36c..4a5bfec9967 100644
--- a/lib/rspec_flaky/listener.rb
+++ b/lib/rspec_flaky/listener.rb
@@ -2,11 +2,15 @@ require 'json'
 
 module RspecFlaky
   class Listener
-    attr_reader :all_flaky_examples, :new_flaky_examples
-
-    def initialize
-      @new_flaky_examples = {}
-      @all_flaky_examples = init_all_flaky_examples
+    # - suite_flaky_examples: contains all the currently tracked flacky example
+    #   for the whole RSpec suite
+    # - flaky_examples: contains the examples detected as flaky during the
+    #   current RSpec run
+    attr_reader :suite_flaky_examples, :flaky_examples
+
+    def initialize(suite_flaky_examples_json = nil)
+      @flaky_examples = FlakyExamplesCollection.new
+      @suite_flaky_examples = init_suite_flaky_examples(suite_flaky_examples_json)
     end
 
     def example_passed(notification)
@@ -14,29 +18,21 @@ module RspecFlaky
 
       return unless current_example.attempts > 1
 
-      flaky_example_hash = all_flaky_examples[current_example.uid]
-
-      all_flaky_examples[current_example.uid] =
-        if flaky_example_hash
-          FlakyExample.new(flaky_example_hash).tap do |ex|
-            ex.last_attempts_count = current_example.attempts
-            ex.flaky_reports += 1
-          end
-        else
-          FlakyExample.new(current_example).tap do |ex|
-            new_flaky_examples[current_example.uid] = ex
-          end
-        end
+      flaky_example = suite_flaky_examples.fetch(current_example.uid) { FlakyExample.new(current_example) }
+      flaky_example.update_flakiness!(last_attempts_count: current_example.attempts)
+
+      flaky_examples[current_example.uid] = flaky_example
     end
 
     def dump_summary(_)
-      write_report_file(all_flaky_examples, all_flaky_examples_report_path)
+      write_report_file(flaky_examples, RspecFlaky::Config.flaky_examples_report_path)
 
+      new_flaky_examples = flaky_examples - suite_flaky_examples
       if new_flaky_examples.any?
         Rails.logger.warn "\nNew flaky examples detected:\n"
-        Rails.logger.warn JSON.pretty_generate(to_report(new_flaky_examples))
+        Rails.logger.warn JSON.pretty_generate(new_flaky_examples.to_report)
 
-        write_report_file(new_flaky_examples, new_flaky_examples_report_path)
+        write_report_file(new_flaky_examples, RspecFlaky::Config.new_flaky_examples_report_path)
       end
     end
 
@@ -46,30 +42,23 @@ module RspecFlaky
 
     private
 
-    def init_all_flaky_examples
-      return {} unless File.exist?(all_flaky_examples_report_path)
+    def init_suite_flaky_examples(suite_flaky_examples_json = nil)
+      unless suite_flaky_examples_json
+        return {} unless File.exist?(RspecFlaky::Config.suite_flaky_examples_report_path)
 
-      all_flaky_examples = JSON.parse(File.read(all_flaky_examples_report_path))
+        suite_flaky_examples_json = File.read(RspecFlaky::Config.suite_flaky_examples_report_path)
+      end
 
-      Hash[(all_flaky_examples || {}).map { |k, ex| [k, FlakyExample.new(ex)] }]
+      FlakyExamplesCollection.from_json(suite_flaky_examples_json)
     end
 
-    def write_report_file(examples, file_path)
-      return unless ENV['FLAKY_RSPEC_GENERATE_REPORT'] == 'true'
+    def write_report_file(examples_collection, file_path)
+      return unless RspecFlaky::Config.generate_report?
 
       report_path_dir = File.dirname(file_path)
       FileUtils.mkdir_p(report_path_dir) unless Dir.exist?(report_path_dir)
-      File.write(file_path, JSON.pretty_generate(to_report(examples)))
-    end
-
-    def all_flaky_examples_report_path
-      @all_flaky_examples_report_path ||= ENV['ALL_FLAKY_RSPEC_REPORT_PATH'] ||
-        Rails.root.join("rspec_flaky/all-report.json")
-    end
 
-    def new_flaky_examples_report_path
-      @new_flaky_examples_report_path ||= ENV['NEW_FLAKY_RSPEC_REPORT_PATH'] ||
-        Rails.root.join("rspec_flaky/new-report.json")
+      File.write(file_path, JSON.pretty_generate(examples_collection.to_report))
     end
   end
 end