From 25feec6c4cfc38def3b7e74f3be734467f5d0a87 Mon Sep 17 00:00:00 2001
From: Andrey Marchenko <andrey.marchenko@datadoghq.com>
Date: Wed, 22 Nov 2023 14:37:17 +0100
Subject: [PATCH] Datadog::CI public methods for test session tracing

---
 lib/datadog/ci.rb           | 53 +++++++++++++++++++++++++++++++++++++
 sig/datadog/ci.rbs          |  6 +++++
 sig/datadog/ci/recorder.rbs |  2 ++
 spec/datadog/ci_spec.rb     | 34 ++++++++++++++++++++++++
 4 files changed, 95 insertions(+)

diff --git a/lib/datadog/ci.rb b/lib/datadog/ci.rb
index 7bfd382e3..96dc93473 100644
--- a/lib/datadog/ci.rb
+++ b/lib/datadog/ci.rb
@@ -10,6 +10,54 @@ module Datadog
   # @public_api
   module CI
     class << self
+      # Return a {Datadog::CI::TestSesstion ci_test_session} that represents the whole test session run.
+      # Raises an error if a session is already active.
+      #
+      #
+      # The {#start_test_session} method is used to mark the start :
+      # ```
+      # Datadog::CI.start_test_session(
+      #   service: "my-web-site-tests",
+      #   tags: { Datadog::CI::Ext::Test::TAG_FRAMEWORK => "my-test-framework" }
+      # )
+      #
+      # # Somewhere else after test run has ended
+      # Datadog::CI.active_test_session.finish
+      # ```
+      #
+      # Remember that calling {Datadog::CI::TestSession#finish} is mandatory.
+      #
+      # @param [String] service_name the service name for this session
+      # @param [Hash<String,String>] tags extra tags which should be added to the test.
+      # @return [Datadog::CI::TestSession] returns the active, running {Datadog::CI::TestSession}.
+      #
+      # @public_api
+      def start_test_session(service_name: nil, tags: {})
+        recorder.start_test_session(service_name: service_name, tags: tags)
+      end
+
+      # The active, unfinished test session span.
+      #
+      # Usage:
+      #
+      # ```
+      # # start a test session
+      # Datadog::CI.start_test_session(
+      #   service: "my-web-site-tests",
+      #   tags: { Datadog::CI::Ext::Test::TAG_FRAMEWORK => "my-test-framework" }
+      # )
+      #
+      # # somewhere else, access the session
+      # test_session = Datadog::CI.active_test_session
+      # test_session.finish
+      # ```
+      #
+      # @return [Datadog::CI::TestSession] the active test session
+      # @return [nil] if no test session is active
+      def active_test_session
+        recorder.active_test_session
+      end
+
       # Return a {Datadog::CI::Test ci_test} that will trace a test called `test_name`.
       # Raises an error if a test is already active.
       #
@@ -189,6 +237,11 @@ def deactivate_test(test)
         recorder.deactivate_test(test)
       end
 
+      # Internal only, to finish a test session use Datadog::CI::TestSession#finish
+      def deactivate_test_session
+        recorder.deactivate_test_session
+      end
+
       private
 
       def components
diff --git a/sig/datadog/ci.rbs b/sig/datadog/ci.rbs
index 9a1b76b46..2340fccdb 100644
--- a/sig/datadog/ci.rbs
+++ b/sig/datadog/ci.rbs
@@ -4,14 +4,20 @@ module Datadog
 
     def self.start_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) -> Datadog::CI::Test
 
+    def self.start_test_session: (?service_name: String?, ?tags: Hash[untyped, untyped]) -> Datadog::CI::TestSession
+
     def self.trace: (String span_type, String span_name, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Span span) -> untyped } -> untyped
 
+    def self.active_test_session: () -> Datadog::CI::TestSession?
+
     def self.active_test: () -> Datadog::CI::Test?
 
     def self.active_span: (String span_type) -> Datadog::CI::Span?
 
     def self.deactivate_test: (Datadog::CI::Test test) -> void
 
+    def self.deactivate_test_session: () -> void
+
     def self.components: () -> Datadog::CI::Configuration::Components
 
     def self.recorder: () -> Datadog::CI::Recorder
diff --git a/sig/datadog/ci/recorder.rbs b/sig/datadog/ci/recorder.rbs
index 8d3d2e097..6bafc8065 100644
--- a/sig/datadog/ci/recorder.rbs
+++ b/sig/datadog/ci/recorder.rbs
@@ -13,6 +13,8 @@ module Datadog
 
       def start_test_session: (?service_name: String?, ?tags: Hash[untyped, untyped]) -> Datadog::CI::TestSession
 
+      def active_test_session: () -> Datadog::CI::TestSession?
+
       def active_test: () -> Datadog::CI::Test?
 
       def active_span: () -> Datadog::CI::Span?
diff --git a/spec/datadog/ci_spec.rb b/spec/datadog/ci_spec.rb
index bee49fbd1..d8a71aa3d 100644
--- a/spec/datadog/ci_spec.rb
+++ b/spec/datadog/ci_spec.rb
@@ -98,4 +98,38 @@
       it { is_expected.to be_nil }
     end
   end
+
+  describe "::start_test_session" do
+    subject(:start_test_session) { described_class.start_test_session }
+
+    let(:ci_test_session) { instance_double(Datadog::CI::TestSession) }
+
+    before do
+      allow(recorder).to receive(:start_test_session).and_return(ci_test_session)
+    end
+
+    it { is_expected.to be(ci_test_session) }
+  end
+
+  describe "::active_test_session" do
+    subject(:active_test_session) { described_class.active_test_session }
+
+    let(:ci_test_session) { instance_double(Datadog::CI::TestSession) }
+
+    before do
+      allow(recorder).to receive(:active_test_session).and_return(ci_test_session)
+    end
+
+    it { is_expected.to be(ci_test_session) }
+  end
+
+  describe "::deactivate_test_session" do
+    subject(:deactivate_test_session) { described_class.deactivate_test_session }
+
+    before do
+      allow(recorder).to receive(:deactivate_test_session)
+    end
+
+    it { is_expected.to be_nil }
+  end
 end