feat: Add MongoengineConnectionField.reference_args

Author: abawchen

graphene_mongo/fields.py

@@ -60,26 +60,34 @@ def model(self):
     @property
     def args(self):
         return to_arguments(
-            self._base_args or OrderedDict(), self.default_filter_args
+            self._base_args or OrderedDict(), dict(self.field_args.items() + self.reference_args.items())
         )
 
     @args.setter
     def args(self, args):
         self._base_args = args
 
     @property
-    def default_filter_args(self):
+    def field_args(self):
         def is_filterable(kv):
             return hasattr(kv[1], '_type') \
                 and callable(getattr(kv[1]._type, '_of_type', None))
 
         return reduce(
             lambda r, kv: r.update(
                 {kv[0]: kv[1]._type._of_type()}) or r if is_filterable(kv) else r,
-            self.fields.items(),
-            {}
+            self.fields.items(), {}
         )
 
+    @property
+    def reference_args(self):
+        def get_reference_field(r, kv):
+            if callable(getattr(kv[1], 'get_type', None)):
+                node = kv[1].get_type()._type._meta
+                r.update({kv[0]: node.fields['id']._type.of_type()})
+            return r
+        return reduce(get_reference_field, self.fields.items(), {})
+
     @property
     def filter_fields(self):
         return self._type._meta.filter_fields
@@ -121,7 +129,7 @@ def get_query(cls, model, info, **args):
             if first is not None:
                 objs = objs[:first]
             if last is not None:
-                # fix for https://github.com/graphql-python/graphene-mongo/issues/20
+                # https://github.com/graphql-python/graphene-mongo/issues/20
                 objs = objs[-(last+1):]
 
         return objs

graphene_mongo/tests/test_fields.py

@@ -0,0 +1,18 @@
+import pytest
+
+from ..fields import MongoengineConnectionField
+from .types import ArticleNode
+
+
+def test_field_args():
+    field = MongoengineConnectionField(ArticleNode)
+
+    field_args = ['id', 'headline', 'pub_date']
+    assert set(field.field_args.keys()) == set(field_args)
+
+    reference_args = ['editor', 'reporter']
+    assert set(field.reference_args.keys()) == set(reference_args)
+
+    default_args = ['after', 'last', 'first', 'before']
+    args = field_args + reference_args + default_args
+    assert set(field.args) == set(args)

graphene_mongo/tests/test_utils.py

@@ -5,14 +5,20 @@
 from ..utils import (
     get_model_fields, is_valid_mongoengine_model
 )
-from .models import Reporter, Child
+from .models import Article, Reporter, Child
+
 
 def test_get_model_fields_no_duplication():
     reporter_fields = get_model_fields(Reporter)
     reporter_name_set = set(reporter_fields)
     assert len(reporter_fields) == len(reporter_name_set)
 
 
+def test_get_model_relation_fields():
+    article_fields = get_model_fields(Article)
+    assert all(field in set(article_fields) for field in ['editor', 'reporter'])
+
+
 def test_get_base_model_fields():
     child_fields = get_model_fields(Child)
     assert all(field in set(child_fields) for field in ['bar', 'baz'])