From 2d24951a0eb89f3482b01c75d8f2aebfbcdf5a34 Mon Sep 17 00:00:00 2001
From: Andrey Marchenko <andrey.marchenko@datadoghq.com>
Date: Wed, 21 Feb 2024 13:14:22 +0100
Subject: [PATCH] sinmplest code coverage with resume/suspend POC

---
 lib/datadog/ci/itr/coverage/collector.rb   | 47 ++++++++++++++++++++++
 lib/datadog/ci/itr/coverage/filter.rb      | 29 +++++++++++++
 lib/datadog/ci/test_visibility/recorder.rb | 25 ++++++++++++
 3 files changed, 101 insertions(+)
 create mode 100644 lib/datadog/ci/itr/coverage/collector.rb
 create mode 100644 lib/datadog/ci/itr/coverage/filter.rb

diff --git a/lib/datadog/ci/itr/coverage/collector.rb b/lib/datadog/ci/itr/coverage/collector.rb
new file mode 100644
index 00000000..ac6ed911
--- /dev/null
+++ b/lib/datadog/ci/itr/coverage/collector.rb
@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require "coverage"
+
+require_relative "filter"
+
+module Datadog
+  module CI
+    module Itr
+      module Coverage
+        class Collector
+          def initialize
+            # Do not run code coverage if someone else is already running it.
+            # It means that user is running the test with coverage and ITR would mess it up.
+            @coverage_supported = !::Coverage.running?
+            # @coverage_supported = false
+          end
+
+          def setup
+            if @coverage_supported
+              p "RUNNING WITH CODE COVERAGE ENABLED!"
+              ::Coverage.setup(lines: true)
+            else
+              p "RUNNING WITH CODE COVERAGE DISABLED!"
+            end
+          end
+
+          def start
+            return unless @coverage_supported
+
+            # if execution is threaded then coverage might already be running
+            ::Coverage.resume unless ::Coverage.running?
+          end
+
+          def stop
+            return nil unless @coverage_supported
+
+            result = ::Coverage.result(stop: false, clear: true)
+            ::Coverage.suspend if ::Coverage.running?
+
+            Filter.call(result)
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/datadog/ci/itr/coverage/filter.rb b/lib/datadog/ci/itr/coverage/filter.rb
new file mode 100644
index 00000000..d81ad969
--- /dev/null
+++ b/lib/datadog/ci/itr/coverage/filter.rb
@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require_relative "../../utils/git"
+
+module Datadog
+  module CI
+    module Itr
+      module Coverage
+        class Filter
+          def self.call(raw_result)
+            new.call(raw_result)
+          end
+
+          def initialize(root: Utils::Git.root)
+            @regex = /\A#{Regexp.escape(root + File::SEPARATOR)}/i.freeze
+          end
+
+          def call(raw_result)
+            return nil if raw_result.nil?
+
+            raw_result.select do |path, coverage|
+              path =~ @regex && coverage[:lines].any? { |count| count && count > 0 }
+            end
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/datadog/ci/test_visibility/recorder.rb b/lib/datadog/ci/test_visibility/recorder.rb
index c274590f..bd36e11e 100644
--- a/lib/datadog/ci/test_visibility/recorder.rb
+++ b/lib/datadog/ci/test_visibility/recorder.rb
@@ -20,6 +20,8 @@
 require_relative "../test_module"
 require_relative "../test_suite"
 
+require_relative "../itr/coverage/collector"
+
 module Datadog
   module CI
     module TestVisibility
@@ -38,6 +40,13 @@ def initialize(
           @local_context = Context::Local.new
           @global_context = Context::Global.new
           @codeowners = codeowners
+
+          @coverage_collector = Itr::Coverage::Collector.new
+          begin
+            @coverage_collector.setup
+          rescue => e
+            Datadog.logger.debug("Coverage collector setup failed: #{e.message}")
+          end
         end
 
         def start_test_session(service: nil, tags: {})
@@ -103,12 +112,15 @@ def trace_test(test_name, test_suite_name, service: nil, tags: {}, &block)
             {resource: test_name, continue_from: Datadog::Tracing::TraceDigest.new}
           )
 
+          @coverage_collector.start
           if block
             start_datadog_tracer_span(test_name, span_options) do |tracer_span|
               test = build_test(tracer_span, tags)
 
               @local_context.activate_test(test) do
                 block.call(test)
+
+                on_test_end(test)
               end
             end
           else
@@ -160,6 +172,8 @@ def active_test_suite(test_suite_name)
         end
 
         def deactivate_test
+          on_test_end(active_test) if active_test
+
           @local_context.deactivate_test
         end
 
@@ -314,6 +328,17 @@ def validate_test_suite_level_visibility_correctness(test)
             end
           end
         end
+
+        # event system?
+        def on_test_end(test)
+          coverage = @coverage_collector.stop
+          if coverage
+            files_covered = coverage.keys.map { |filename| Utils::Git.relative_to_root(filename) }
+            test.set_tag("_test.coverage", files_covered)
+
+            # p test.get_tag("_test.coverage")
+          end
+        end
       end
     end
   end