extract set_session_context and set_module_context in Recorder

DataDog · Nov 30, 2023 · b48270f · b48270f
1 parent 64ad04f
commit b48270f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 21 deletions.
diff --git a/lib/datadog/ci/recorder.rb b/lib/datadog/ci/recorder.rb
@@ -21,6 +21,7 @@
 module Datadog
   module CI
     # Common behavior for CI tests
+    # Note: this class has too many responsibilities and should be split into multiple classes
     class Recorder
       attr_reader :environment_tags, :test_suite_level_visibility_enabled, :enabled
 
@@ -39,8 +40,7 @@ def start_test_session(service_name: nil, tags: {})
         tracer_span = start_datadog_tracer_span(
           "test.session", build_span_options(service_name, Ext::AppTypes::TYPE_TEST_SESSION)
         )
-
-        tags[Ext::Test::TAG_TEST_SESSION_ID] = tracer_span.id
+        set_session_context(tags, tracer_span)
 
         test_session = build_test_session(tracer_span, tags)
         @global_context.activate_test_session!(test_session)
@@ -52,18 +52,12 @@ def start_test_module(test_module_name, service_name: nil, tags: {})
         return skip_tracing unless test_suite_level_visibility_enabled
 
         tags = tags_with_inherited_globals(tags)
-
-        test_session = active_test_session
-        if test_session
-          tags[Ext::Test::TAG_TEST_SESSION_ID] = test_session.id
-        end
+        set_session_context(tags)
 
         tracer_span = start_datadog_tracer_span(
           test_module_name, build_span_options(service_name, Ext::AppTypes::TYPE_TEST_MODULE)
         )
-
-        tags[Ext::Test::TAG_TEST_MODULE_ID] = tracer_span.id
-        tags[Ext::Test::TAG_MODULE] = tracer_span.name
+        set_module_context(tags, tracer_span)
 
         test_module = build_test_module(tracer_span, tags)
         @global_context.activate_test_module!(test_module)
@@ -76,17 +70,8 @@ def trace_test(test_name, service_name: nil, operation_name: "test", tags: {}, &
         return skip_tracing(block) unless enabled
 
         tags = tags_with_inherited_globals(tags)
-
-        test_session = active_test_session
-        if test_session
-          tags[Ext::Test::TAG_TEST_SESSION_ID] = test_session.id
-        end
-
-        test_module = active_test_module
-        if test_module
-          tags[Ext::Test::TAG_TEST_MODULE_ID] = test_module.id
-          tags[Ext::Test::TAG_MODULE] = test_module.name
-        end
+        set_session_context(tags)
+        set_module_context(tags)
 
         tags[Ext::Test::TAG_NAME] = test_name
 
@@ -221,6 +206,19 @@ def set_initial_tags(ci_span, tags)
         ci_span.set_tags(environment_tags)
       end
 
+      def set_session_context(tags, test_session = nil)
+        test_session ||= active_test_session
+        tags[Ext::Test::TAG_TEST_SESSION_ID] = test_session.id if test_session
+      end
+
+      def set_module_context(tags, test_module = nil)
+        test_module ||= active_test_module
+        if test_module
+          tags[Ext::Test::TAG_TEST_MODULE_ID] = test_module.id
+          tags[Ext::Test::TAG_MODULE] = test_module.name
+        end
+      end
+
       def start_datadog_tracer_span(span_name, span_options, &block)
         if block
           Datadog::Tracing.trace(span_name, **span_options) do |tracer_span, trace|

diff --git a/sig/datadog/ci/recorder.rbs b/sig/datadog/ci/recorder.rbs
@@ -56,6 +56,11 @@ module Datadog
 
       def set_initial_tags: (Datadog::CI::Span ci_span, Hash[untyped, untyped] tags) -> void
 
+      # the type (Datadog::CI::TestSession | Datadog::Tracing::SpanOperation) screams of wrong/mising abstraction
+      def set_session_context: (Hash[untyped, untyped] tags, ?Datadog::CI::TestSession | Datadog::Tracing::SpanOperation? test_session) -> void
+
+      def set_module_context: (Hash[untyped, untyped] tags, ?Datadog::CI::TestModule | Datadog::Tracing::SpanOperation? test_module) -> void
+
       def null_span: () -> Datadog::CI::Span
 
       def skip_tracing: (?untyped block) -> untyped