fix(tracer): add support for multiplex (#243)

* fix(tracer): add support for multiplex

* style(tracer): make rubocop happier

---------

Co-authored-by: ashleyyy <[email protected]>

Author: ashleyyy

lib/apollo-federation/tracing/tracer.rb

@@ -83,16 +83,26 @@ def self.execute_query_lazy(data, &block)
         result = block.call
 
         query = data.fetch(:query)
-        return result unless query.context && query.context[:tracing_enabled]
+        multiplex = data.fetch(:multiplex)
+
+        if query
+          record_trace_end_time(query)
+        elsif multiplex
+          multiplex.queries.each { |q| record_trace_end_time(q) }
+        end
+
+        result
+      end
+
+      def self.record_trace_end_time(query)
+        return unless query.context && query.context[:tracing_enabled]
 
         trace = query.context.namespace(ApolloFederation::Tracing::KEY)
 
         trace.merge!(
           end_time: Time.now.utc,
           end_time_nanos: Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond),
         )
-
-        result
       end
 
       # Step 2:

spec/apollo-federation/tracing_spec.rb

@@ -49,6 +49,41 @@ def test
         result = schema.execute('{ test }', context: { tracing_enabled: true, debug_tracing: true })
         expect(result[:extensions][:ftv1_debug]).not_to be_nil
       end
+
+      context 'with a multiplex query' do
+        it 'does not add tracing extension by default' do
+          result = schema.multiplex(
+            [
+              { query: '{ test }' },
+              { query: '{ test }' },
+              { query: '{ test }' },
+            ],
+          ).first
+          expect(result[:extensions]).to be_nil
+        end
+
+        it 'adds the extensions.ftv1 when the context has tracing_enabled: true' do
+          result = schema.multiplex(
+            [
+              { query: '{ test }', context: { tracing_enabled: true } },
+              { query: '{ test }', context: { tracing_enabled: true } },
+              { query: '{ test }', context: { tracing_enabled: true } },
+            ],
+          ).first
+          expect(result[:extensions][:ftv1]).not_to be_nil
+        end
+
+        it 'adds debugging info when the context has debug_tracing: true' do
+          result = schema.multiplex(
+            [
+              { query: '{ test }', context: { tracing_enabled: true, debug_tracing: true } },
+              { query: '{ test }', context: { tracing_enabled: true, debug_tracing: true } },
+              { query: '{ test }', context: { tracing_enabled: true, debug_tracing: true } },
+            ],
+          ).first
+          expect(result[:extensions][:ftv1_debug]).not_to be_nil
+        end
+      end
     end
 
     def trace(query)