fix: ensure query object is traversed

lib/apollo-federation/schema.rb

@@ -42,18 +42,15 @@ def link_namespace
       end
 
       def query(new_query_object = nil)
-        if new_query_object
-          @orig_query_object = new_query_object
-        else
-          if !@federation_query_object
-            @federation_query_object = federation_query(original_query)
-            @federation_query_object.define_entities_field(schema_entities)
+        return super if new_query_object.nil? && @query_object
 
-            super(@federation_query_object)
-          end
+        @orig_query_object = new_query_object
+        federation_query_object = federation_query(original_query)
+        federation_query_object.define_entities_field(schema_entities)
 
-          super
-        end
+        super(federation_query_object)
+
+        federation_query_object
       end
 
       private

spec/apollo-federation/schema_spec.rb

@@ -3,6 +3,7 @@
 require 'spec_helper'
 require 'graphql'
 require 'apollo-federation/schema'
+require 'apollo-federation/object'
 
 RSpec.describe ApolloFederation::Schema do
   describe '.federation_version' do
@@ -97,4 +98,22 @@
       expect(schema.federation_2?).to be(true)
     end
   end
+
+  describe '.query' do
+    it 'traverses the query type' do
+      cat_type = Class.new(GraphQL::Schema::Object) do
+        graphql_name 'Cat'
+      end
+      query_type = Class.new(GraphQL::Schema::Object) do
+        graphql_name 'Query'
+        field :cat, cat_type, null: false
+      end
+      schema = Class.new(GraphQL::Schema) do
+        include ApolloFederation::Schema
+        query query_type
+      end
+
+      expect(schema.get_type('Cat')).to eq(cat_type)
+    end
+  end
 end