Address PR review feedback on GraphQL::Schema::Type helpers

- `source_type`: restore `&.` safe navigation since `@object_runtime_metadata`
  can be nil for non-object types (scalars, enums); return `nil` instead of
  `self` for non-derived types; update `search_index_definitions` to branch on
  `source_type` truthy rather than `indexed_aggregation?`, and update comment
  to describe the general case with indexed aggregation as a prime example
- `subtypes`: return a `Set` instead of `Array`; inline `subtypes` in
  `non_subtypes_in_shared_index` rather than caching in a local variable
- Tests: use `before(:context)` shared schema in `#source_type` group; add
  scalar and enum nil cases; use `contain_exactly` for precise set assertions;
  consolidate `not_to include` calls; add coverage for `t == self` guard

Generated with Claude Code

Author: marcdaniels-toast

elasticgraph-graphql/lib/elastic_graph/graphql/schema/type.rb

@@ -61,12 +61,12 @@ def name
         # List of index definitions that should be searched for this type.
         def search_index_definitions
           @search_index_definitions ||=
-            if indexed_aggregation?
-              # For an indexed aggregation, we just delegate to its source type. This works better than
-              # dumping index definitions in the runtime metadata of the indexed aggregation type itself
-              # because of abstract (interface/union) types. The source document type handles that (since
-              # there is a supertype/subtype relationship on the document types) but that relationship
-              # does not exist on the indexed aggregation.
+            if (st = source_type)
+              # When a type has a source type (a prime example being indexed aggregations), we delegate
+              # to the source type. This works better than dumping index definitions in the runtime metadata
+              # of the derived type itself because of abstract (interface/union) types. The source document
+              # type handles that (since there is a supertype/subtype relationship on the document types)
+              # but that relationship does not exist on derived types.
               #
               # For example, assume we have these indexed document types:
               # - type Person {}
@@ -76,7 +76,7 @@ def search_index_definitions
               # We can go from `Inventor` to its subtypes to find the search indexes. However, `InventorAggregation`
               # is NOT a union of `PersonAggregation` and `CompanyAggregation`, so we can't do the same thing on the
               # indexed aggregation types. Delegating to the source type solves this case.
-              source_type.search_index_definitions
+              st.search_index_definitions
             else
               @index_definitions.union(subtypes.flat_map(&:search_index_definitions))
             end
@@ -123,18 +123,15 @@ def subtypes
           @subtypes ||= @schema
             .graphql_schema
             .possible_types(graphql_type, visibility_profile: :boot)
-            .map { |t| @schema.type_from(t) } - [self]
+            .map { |t| @schema.type_from(t) }
+            .reject { |t| t == self }
+            .to_set
         end
 
         # For derived types (e.g. indexed aggregations), returns the underlying source document type.
-        # For non-derived types, returns `self`.
+        # Returns `nil` for non-derived types.
         def source_type
-          @source_type ||=
-            if (st = @object_runtime_metadata.source_type)
-              @schema.type_named(st)
-            else
-              self
-            end
+          @source_type ||= @object_runtime_metadata&.source_type&.then { |st| @schema.type_named(st) }
         end
 
         # Returns the set of concrete indexed document types that share any of this type's search
@@ -145,13 +142,11 @@ def source_type
         # with a concrete `__typename`. When filtering by `__typename`, only concrete types are
         # relevant.
         def non_subtypes_in_shared_index
-          @non_subtypes_in_shared_index ||= begin
-            all_subtypes = subtypes.to_set # all concrete subtypes at any depth
+          @non_subtypes_in_shared_index ||=
             search_index_definitions
               .flat_map { |index_def| @schema.document_types_stored_in(index_def.name).to_a }
-              .reject { |t| t == self || all_subtypes.include?(t) || t.abstract? }
+              .reject { |t| t == self || subtypes.include?(t) || t.abstract? }
               .to_set
-          end
         end
 
         def field_named(field_name)

elasticgraph-graphql/sig/elastic_graph/graphql/schema/type.rbs

@@ -9,8 +9,8 @@ module ElasticGraph
         attr_reader graphql_type: ::GraphQL::Schema::_Type
         attr_reader grouping_missing_value_placeholder: ::String? | ::Numeric?
         def search_index_definitions: () -> ::Array[DatastoreCore::_IndexDefinition]
-        def source_type: () -> Type
-        def subtypes: () -> ::Array[Type]
+        def source_type: () -> Type?
+        def subtypes: () -> ::Set[Type]
         def non_subtypes_in_shared_index: () -> ::Set[Type]
         def unwrap_fully: () -> Type
         def field_named: (::String) -> Field

elasticgraph-graphql/spec/unit/elastic_graph/graphql/schema/type_spec.rb

@@ -604,14 +604,14 @@ def type_for(field_name)
             end
           end
 
-          it "returns [] for object types" do
+          it "returns an empty set for object types" do
             type = schema.type_named("Size")
-            expect(type.subtypes).to eq []
+            expect(type.subtypes).to be_empty
           end
 
-          it "returns [] for scalar types" do
+          it "returns an empty set for scalar types" do
             type = schema.type_named("Int")
-            expect(type.subtypes).to eq []
+            expect(type.subtypes).to be_empty
           end
 
           it "returns the subtypes of a union" do
@@ -759,28 +759,37 @@ def search_index_definitions_from(type_name: "TheType")
         end
 
         describe "#source_type" do
-          it "returns the underlying document type for an indexed aggregation type" do
-            schema = define_schema do |s|
+          attr_reader :schema
+
+          before(:context) do
+            @schema = define_schema(clients_by_name: {}) do |s|
+              s.enum_type "Status" do |t|
+                t.value "ACTIVE"
+                t.value "INACTIVE"
+              end
+
               s.object_type "Thing" do |t|
                 t.field "id", "ID!"
+                t.field "status", "Status"
                 t.index "things"
               end
             end
+          end
 
-            aggregation_type = schema.type_named("ThingAggregation")
-            expect(aggregation_type.source_type).to be schema.type_named("Thing")
+          it "returns the underlying document type for an indexed aggregation type" do
+            expect(schema.type_named("ThingAggregation").source_type).to be schema.type_named("Thing")
           end
 
-          it "returns self for a non-derived type" do
-            schema = define_schema do |s|
-              s.object_type "Thing" do |t|
-                t.field "id", "ID!"
-                t.index "things"
-              end
-            end
+          it "returns nil for a non-derived object type" do
+            expect(schema.type_named("Thing").source_type).to be_nil
+          end
 
-            thing_type = schema.type_named("Thing")
-            expect(thing_type.source_type).to be thing_type
+          it "returns nil for a scalar type" do
+            expect(schema.type_named("Int").source_type).to be_nil
+          end
+
+          it "returns nil for an enum type" do
+            expect(schema.type_named("Status").source_type).to be_nil
           end
         end
 
@@ -821,17 +830,14 @@ def search_index_definitions_from(type_name: "TheType")
             end
           end
 
-          it "excludes the type itself and its subtypes" do
-            store = schema.type_named("Store")
-
-            result = store.non_subtypes_in_shared_index
-            expect(result).not_to include(store)
-            expect(result).not_to include(schema.type_named("OnlineStore"))
-            expect(result).not_to include(schema.type_named("PhysicalStore"))
+          it "excludes the type itself and its subtypes, returning only concrete sibling types in the shared index" do
+            expect(schema.type_named("Store").non_subtypes_in_shared_index).to contain_exactly(schema.type_named("Wholesaler"))
           end
 
-          it "includes concrete sibling types that share the same index" do
-            expect(schema.type_named("Store").non_subtypes_in_shared_index).to include(schema.type_named("Wholesaler"))
+          it "excludes the type itself even when the type is a concrete type in a shared index" do
+            # Wholesaler is a concrete type stored in the "channels" index, so it appears in
+            # document_types_stored_in("channels"). Without the `t == self` guard it would include itself.
+            expect(schema.type_named("Wholesaler").non_subtypes_in_shared_index).to contain_exactly(schema.type_named("OnlineStore"))
           end
 
           it "returns an empty set when all types sharing its indexes are subtypes" do