fix: handle force_flush for rake task with arguments (#1412)

Co-authored-by: Ariel Valentin <[email protected]>
Co-authored-by: Kayla Reopelle <[email protected]>

Author: 3 people

instrumentation/rake/lib/opentelemetry/instrumentation/rake/patches/task.rb

@@ -29,7 +29,11 @@ def tracer
           end
 
           def force_flush
-            OpenTelemetry.tracer_provider.force_flush if ::Rake.application.top_level_tasks.include?(name)
+            top_level_task_names = ::Rake.application.top_level_tasks.map { |t| t.split('[').first }
+
+            if top_level_task_names.include?(name)
+              OpenTelemetry.tracer_provider.force_flush
+            end
           end
         end
       end

instrumentation/rake/test/opentelemetry/instrumentation/rake/patches/task_test.rb

@@ -56,6 +56,25 @@
       _(execute_span.parent_span_id).must_equal(invoke_span.span_id)
     end
 
+    describe 'with a task argument' do
+      it 'should call force_flush on OpenTelemetry.tracer_provider' do
+        mock = Minitest::Mock.new
+        mock.expect(:force_flush, nil)
+        mock.expect(:force_flush, nil)
+
+        Rake::Task.define_task("#{task_name}[:arg]")
+        task_string = "#{task_name}[test_arg]"
+
+        Rake.application.instance_eval { @top_level_tasks = [task_string] }
+
+        OpenTelemetry.stub(:tracer_provider, mock) do
+          Rake.application.invoke_task(task_string)
+        end
+
+        mock.verify
+      end
+    end
+
     describe 'with a prerequisite task' do
       before do
         Rake::Task.define_task(prerequisite_task_name)