diff --git a/lib/graphql/schema.rb b/lib/graphql/schema.rb index 24cb8e9b6f..356335ea85 100644 --- a/lib/graphql/schema.rb +++ b/lib/graphql/schema.rb @@ -475,8 +475,10 @@ def subscription(new_subscription_object = nil, &lazy_load_block) raise GraphQL::Error, "Second definition of `subscription(...)` (#{dup_defn.inspect}) is invalid, already configured with #{@subscription_object.inspect}" elsif use_schema_subset? @subscription_object = block_given? ? lazy_load_block : new_subscription_object + add_subscription_extension_if_necessary else @subscription_object = new_subscription_object || lazy_load_block.call + add_subscription_extension_if_necessary add_type_and_traverse(@subscription_object, root: true) end nil diff --git a/spec/graphql/schema_spec.rb b/spec/graphql/schema_spec.rb index d6f07f1a8f..5a0c8cb307 100644 --- a/spec/graphql/schema_spec.rb +++ b/spec/graphql/schema_spec.rb @@ -495,5 +495,23 @@ class QueryRequiredSchema < GraphQL::Schema assert_equal [:query, :mutation], calls assert_equal "Subscription", schema.subscription.graphql_name assert_equal [:query, :mutation, :subscription], calls + assert schema.instance_variable_get(:@subscription_extension_added) + end + + it "adds the subscription extension if subscription(...) is called second" do + schema = Class.new(GraphQL::Schema) do + use GraphQL::Subscriptions + subscription(Class.new(GraphQL::Schema::Object) { graphql_name("Subscription") }) + end + assert schema.subscription + assert schema.instance_variable_get(:@subscription_extension_added) + + schema2 = Class.new(GraphQL::Schema) do + self.use_schema_subset = true + use GraphQL::Subscriptions + subscription(Class.new(GraphQL::Schema::Object) { graphql_name("Subscription") }) + end + assert schema2.subscription + assert schema2.instance_variable_get(:@subscription_extension_added) end end