fix[union-converter]: get_queried_union_types can now handle union fragments

Author: mak626

graphene_mongo/converter.py

@@ -186,7 +186,9 @@ def reference_resolver(root, *args, **kwargs):
                     return None
 
                 choice_to_resolve = dict()
-                querying_union_types = get_queried_union_types(args[0])
+                querying_union_types = get_queried_union_types(
+                    info=args[0], valid_gql_types=registry._registry_string_map.keys()
+                )
                 to_resolve_models = dict()
                 for each, queried_fields in querying_union_types.items():
                     to_resolve_models[registry._registry_string_map[each]] = queried_fields
@@ -263,7 +265,9 @@ async def reference_resolver_async(root, *args, **kwargs):
                     return None
 
                 choice_to_resolve = dict()
-                querying_union_types = get_queried_union_types(args[0])
+                querying_union_types = get_queried_union_types(
+                    info=args[0], valid_gql_types=registry._registry_async_string_map.keys()
+                )
                 to_resolve_models = dict()
                 for each, queried_fields in querying_union_types.items():
                     to_resolve_models[registry._registry_async_string_map[each]] = queried_fields

graphene_mongo/utils.py

@@ -203,11 +203,12 @@ def get_query_fields(info):
     return query
 
 
-def get_queried_union_types(info):
+def get_queried_union_types(info, valid_gql_types):
     """A convenience function to get queried union types with its fields
 
     Args:
         info (ResolveInfo)
+        valid_gql_types (dict_keys)
 
     Returns:
         dict[union_type_name, queried_fields(dict)]
@@ -227,9 +228,16 @@ def get_queried_union_types(info):
         for leaf in selection_set.selections:
             if leaf.kind == "fragment_spread":
                 fragment_name = fragments[leaf.name.value].type_condition.name.value
-                fragments_queries[fragment_name] = collect_query_fields(
+                sub_query_fields = collect_query_fields(
                     fragments[leaf.name.value], fragments, variables
                 )
+                if fragment_name not in valid_gql_types:
+                    # This is done to avoid UnionFragments coming in fragments_queries as
+                    # we actually need its children types and not the UnionFragments itself
+                    fragments_queries.update(sub_query_fields)
+                    fragments_queries.pop('__typename', None)  # cannot resolve __typename for a union type
+                else:
+                    fragments_queries[fragment_name] = sub_query_fields
             elif leaf.kind == "inline_fragment":
                 fragment_name = leaf.type_condition.name.value
                 fragments_queries[fragment_name] = collect_query_fields(leaf, fragments, variables)