Fix subscription setup

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

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