From d9aac571cbd78852701c3e9c9dc6757ce27bd58d Mon Sep 17 00:00:00 2001
From: Robert Mosolgo <rdmosolgo@gmail.com>
Date: Thu, 8 Aug 2024 11:00:45 -0400
Subject: [PATCH] Add tests for root types from blocks

---
 lib/graphql/schema.rb              |  5 ++++-
 spec/graphql/schema_spec.rb        | 20 ++++++++++++++++++++
 spec/graphql/subscriptions_spec.rb |  4 ++--
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/lib/graphql/schema.rb b/lib/graphql/schema.rb
index d86c9650cd..24cb8e9b6f 100644
--- a/lib/graphql/schema.rb
+++ b/lib/graphql/schema.rb
@@ -482,6 +482,8 @@ def subscription(new_subscription_object = nil, &lazy_load_block)
           nil
         elsif @subscription_object.is_a?(Proc)
           @subscription_object = @subscription_object.call
+          add_subscription_extension_if_necessary
+          @subscription_object
         else
           @subscription_object || find_inherited_value(:subscription)
         end
@@ -1387,7 +1389,8 @@ def instrumenters
 
       # @api private
       def add_subscription_extension_if_necessary
-        if !defined?(@subscription_extension_added) && subscription && self.subscriptions
+        # TODO: when there's a proper API for extending root types, migrat this to use it.
+        if !defined?(@subscription_extension_added) && @subscription_object.is_a?(Class) && self.subscriptions
           @subscription_extension_added = true
           subscription.all_field_definitions.each do |field|
             if !field.extensions.any? { |ext| ext.is_a?(Subscriptions::DefaultSubscriptionResolveExtension) }
diff --git a/spec/graphql/schema_spec.rb b/spec/graphql/schema_spec.rb
index f7a9519f38..d6f07f1a8f 100644
--- a/spec/graphql/schema_spec.rb
+++ b/spec/graphql/schema_spec.rb
@@ -476,4 +476,24 @@ class QueryRequiredSchema < GraphQL::Schema
   it "starts with no references_to" do
     assert_equal({}, GraphQL::Schema.references_to)
   end
+
+  it "defers root type blocks until those types are used" do
+    calls = []
+    schema = Class.new(GraphQL::Schema) do
+      self.use_schema_subset = true
+      query { calls << :query; Class.new(GraphQL::Schema::Object) { graphql_name("Query") } }
+      mutation { calls << :mutation; Class.new(GraphQL::Schema::Object) { graphql_name("Mutation") } }
+      subscription { calls << :subscription; Class.new(GraphQL::Schema::Object) { graphql_name("Subscription") } }
+      # Test this because it tries to modify `subscription` -- currently hardcoded in Schema.add_subscription_extension_if_necessary
+      use GraphQL::Subscriptions
+    end
+
+    assert_equal [], calls
+    assert_equal "Query", schema.query.graphql_name
+    assert_equal [:query], calls
+    assert_equal "Mutation", schema.mutation.graphql_name
+    assert_equal [:query, :mutation], calls
+    assert_equal "Subscription", schema.subscription.graphql_name
+    assert_equal [:query, :mutation, :subscription], calls
+  end
 end
diff --git a/spec/graphql/subscriptions_spec.rb b/spec/graphql/subscriptions_spec.rb
index ef44a83d01..1abdf08dc8 100644
--- a/spec/graphql/subscriptions_spec.rb
+++ b/spec/graphql/subscriptions_spec.rb
@@ -205,8 +205,8 @@ class Query < GraphQL::Schema::Object
   end
 
   class Schema < GraphQL::Schema
-    query(Query)
-    subscription(Subscription)
+    query { Query }
+    subscription { Subscription }
     use InMemoryBackend::Subscriptions, extra: 123
     max_complexity(InMemoryBackend::MAX_COMPLEXITY)
   end