Python mongoengine.Document() Examples

def test_MapField(self):
        class MappedDoc(me.EmbeddedDocument):
            field = me.StringField()
        class Doc(me.Document):
            id = me.IntField(primary_key=True, default=1)
            map = me.MapField(me.EmbeddedDocumentField(MappedDoc))
            str = me.MapField(me.StringField())
        fields_ = fields_for_model(Doc)
        assert type(fields_['map']) is fields.Map
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        doc = Doc(map={'a': MappedDoc(field='A'), 'b': MappedDoc(field='B')},
                  str={'a': 'aaa', 'b': 'bbbb'}).save()
        dump = DocSchema().dump(doc)
        assert not dump.errors
        assert dump.data == {'map': {'a': {'field': 'A'}, 'b': {'field': 'B'}},
                             'str': {'a': 'aaa', 'b': 'bbbb'}, 'id': 1}
        # Try the load
        load = DocSchema().load(dump.data)
        assert not load.errors
        assert load.data.map == doc.map 

def test_invalid_state_transition():
    @acts_as_state_machine
    class Person(mongoengine.Document):
        name = mongoengine.StringField(default='Billy')

        sleeping = State(initial=True)
        running = State()
        cleaning = State()

        run = Event(from_states=sleeping, to_state=running)
        cleanup = Event(from_states=running, to_state=cleaning)
        sleep = Event(from_states=(running, cleaning), to_state=sleeping)

    establish_mongo_connection()
    person = Person()
    person.save()
    assert person.is_sleeping

    #should raise an invalid state exception
    with assert_raises(InvalidStateTransition):
        person.sleep() 

def test_mapping(self):
        class FileDoc(Document):
            f = fields.FileField(collection_name='files')
            i = fields.ImageField(collection_name='images')

        class TestSerializer(DocumentSerializer):
            class Meta:
                model = FileDoc
                fields = '__all__'

        expected = dedent("""
            TestSerializer():
                id = ObjectIdField(read_only=True)
                f = FileField(required=False)
                i = ImageField(required=False)
        """)

        assert repr(TestSerializer()) == expected 

def test_default(self):
        def generate_default_value():
            return 'default_generated_value'
        class Doc(me.Document):
            field_with_default = me.StringField(default='default_value')
            field_required_with_default = me.StringField(required=True,
                default=generate_default_value)
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        # Make sure default doesn't shadow given values
        doc, errors = DocSchema().load({'field_with_default': 'custom_value',
                                        'field_required_with_default': 'custom_value'})
        assert not errors
        assert doc.field_with_default == 'custom_value'
        assert doc.field_required_with_default == 'custom_value'
        # Now use defaults
        doc, errors = DocSchema().load({})
        assert not errors
        assert doc.field_with_default == 'default_value'
        assert doc.field_required_with_default == 'default_generated_value' 

def test_required(self):
        class Doc(me.Document):
            field_not_required = me.StringField()
            field_required = me.StringField(required=True)
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        doc, errors = DocSchema().load({'field_not_required': 'bad_doc'})
        assert errors
        # Now provide the required field
        doc, errors = DocSchema().load({'field_required': 'good_doc'})
        assert not errors
        assert doc.field_not_required is None
        assert doc.field_required == 'good_doc'
        # Update should not take care of the required fields
        doc, errors = DocSchema().update(doc, {'field_not_required': 'good_doc'})
        assert not errors
        assert doc.field_required == 'good_doc'
        assert doc.field_not_required == 'good_doc' 

def test_PointField(self):
        class Doc(me.Document):
            point = me.PointField()
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        doc = Doc(point={ 'type': 'Point', 'coordinates': [10, 20] })
        dump = DocSchema().dump(doc)
        assert not dump.errors
        assert dump.data['point'] == { 'x': 10, 'y': 20 }
        load = DocSchema().load(dump.data)
        assert not load.errors
        assert load.data.point == { 'type': 'Point', 'coordinates': [10, 20] }
        # Deserialize Point with coordinates passed as string
        data = {'point': { 'x': '10', 'y': '20' }}
        load = DocSchema().load(data)
        assert not load.errors
        assert load.data.point == { 'type': 'Point', 'coordinates': [10, 20] }
        # Try to load invalid coordinates
        data = {'point': { 'x': '10', 'y': '20foo' }}
        load = DocSchema().load(data)
        assert 'point' in load.errors 

def test_DynamicField(self):
        class Doc(me.Document):
            dynamic = me.DynamicField()
        fields_ = fields_for_model(Doc)
        assert type(fields_['dynamic']) is fields.Raw
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        data = {
            'int_1': 1,
            'nested_2': {
                'sub_int_1': 42,
                'sub_list_2': []
            },
            'list_3': ['a', 'b', 'c']
        }
        doc = Doc(dynamic=data)
        dump = DocSchema().dump(doc)
        assert not dump.errors
        assert dump.data == {'dynamic': data}
        load = DocSchema().load(dump.data)
        assert not load.errors
        assert load.data.dynamic == data 

def test_ListSpecialField(self):
        class NestedDoc(me.EmbeddedDocument):
            field = me.StringField()
        class Doc(me.Document):
            list = me.ListField(me.EmbeddedDocumentField(NestedDoc))
        fields_ = fields_for_model(Doc)
        assert type(fields_['list']) is fields.List
        assert type(fields_['list'].container) is fields.Nested
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        list_ = [{'field': 'A'}, {'field': 'B'}, {'field': 'C'}]
        doc = Doc(list=list_)
        dump = DocSchema().dump(doc)
        assert not dump.errors
        assert dump.data == {'list': list_}
        load = DocSchema().load(dump.data)
        assert not load.errors
        for i, elem in enumerate(list_):
            assert load.data.list[i].field == elem['field'] 

def test_ListField(self):
        class Doc(me.Document):
            list = me.ListField(me.StringField())
        fields_ = fields_for_model(Doc)
        assert type(fields_['list']) is fields.List
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        list_ = ['A', 'B', 'C']
        doc = Doc(list=list_)
        dump = DocSchema().dump(doc)
        assert not dump.errors
        assert dump.data == {'list': list_}
        load = DocSchema().load(dump.data)
        assert not load.errors
        assert load.data.list == list_ 

def test_FileField(self):
        class File(me.Document):
            name = me.StringField(primary_key=True)
            file = me.FileField()
        class FileSchema(ModelSchema):
            class Meta:
                model = File
        doc = File(name='test_file')
        data = b'1234567890' * 10
        doc.file.put(data, content_type='application/octet-stream')
        dump = FileSchema().dump(doc)
        assert not dump.errors
        assert dump.data == {'name': 'test_file'}
        # Should not be able to load the file
        load = FileSchema().load({'name': 'bad_load', 'file': b'12345'})
        assert not load.data.file 

def __init__(self, items, links=None):  # , session):
        self.forest = {}  # id -> (event/result, parent id, rank)
        """
        :type items: List[Document]
        """
        if links is None:
            links = lambda x: x.links
        self._get_neighbors = links
        # self.session = session

        for item in items:
            self.forest[item.id] = (item, item.id, 0)

        # generate sets
        for item in items:
            for neighbor in links(item):
                self.union(item.id, neighbor.id)

        # for key, (item, item_id, index) in self.forest.items():
        #     for neighbor in self._get_neighbors(item):
        #        self.union(key, neighbor.id)

    # recursively search until you find a parent who is itself 

def test_excluding_declared(self):
        class MockModel(Document):
            foo = fields.StringField()
            bar = fields.StringField()
            baz = fields.StringField()

        class BaseFS(ModelFilterset):
            class Meta:
                model = MockModel
            foo = filters.CharFilter()
            bar = filters.CharFilter()

        class TestFS(BaseFS):
            class Meta:
                model = MockModel
                exclude = ('bar', 'baz')
            baz = filters.CharFilter()
            quz = filters.CharFilter()

        fs = TestFS()
        self.assertEqual(set(fs.filters.keys()), set(['id', 'foo', 'quz'])) 

def test_before_callback_blocking_transition():
    @acts_as_state_machine
    class Runner(mongoengine.Document):
        name = mongoengine.StringField(default='Billy')

        sleeping = State(initial=True)
        running = State()
        cleaning = State()

        run = Event(from_states=sleeping, to_state=running)
        cleanup = Event(from_states=running, to_state=cleaning)
        sleep = Event(from_states=(running, cleaning), to_state=sleeping)

        @before('run')
        def check_sneakers(self):
            return False

    establish_mongo_connection()
    runner = Runner()
    runner.save()
    assert runner.is_sleeping
    runner.run()
    assert runner.is_sleeping
    assert not runner.is_running 

def test_model_subclassed(self):
        """ can apply filters for base model to queryset of derived """
        class FooDoc(Document):
            meta = { 'allow_inheritance': True}
            foo = fields.StringField()

        class BarDoc(FooDoc):
            bar = fields.StringField()

        class FooFilter(ModelFilterset):
            class Meta:
                model = FooDoc

        class BarView(ListAPIView):
            filter_backends = (MongoFilterBackend,)
            filter_class = FooFilter
            serializer_class = mock.Mock()
            queryset = BarDoc.objects

        BarView.as_view()(APIRequestFactory().get("/?foo=Foo")) 

def test_model_mismatch(self):
        """ cannot apply filters for one model to queryset from another """
        class FooDoc(Document):
            foo = fields.StringField()

        class BarDoc(Document):
            bar = fields.StringField()

        class FooFilter(ModelFilterset):
            class Meta:
                model = FooDoc

        class BarView(ListAPIView):
            filter_backends = (MongoFilterBackend,)
            filter_class = FooFilter
            serializer_class = mock.Mock()
            queryset = BarDoc.objects

        with self.assertRaises(TypeError):
            BarView.as_view()(APIRequestFactory().get("/?foo=Foo")) 

def test_skip_none_field(self):
        class Doc(me.Document):
            field_not_empty = me.StringField(default='value')
            field_empty = me.StringField()
            list_empty = me.ListField(me.StringField())
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        doc = Doc()
        dump = DocSchema().dump(doc)
        assert not dump.errors
        assert dump.data == {'field_not_empty': 'value'} 

def save_subdocs(self, validated_data):
        doc = validated_data['ref']
        if isinstance(doc, Document):
            doc.save() 

def to_representation(self, value):
        assert isinstance(value, (Document, DBRef))
        if isinstance(value, Document):
            doc_id = value.id
            doc_cls = value.__class__.__name__
        if isinstance(value, DBRef):  # hard case
            doc_id = value.id
            doc_collection = value.collection
            class_match = [k for k, v in _document_registry.items() if v._get_collection_name() == doc_collection]
            if len(class_match) != 1:
                self.fail('unmapped_collection', collection=doc_collection)
            doc_cls = class_match[0]
        return {'_cls': doc_cls, '_id': self.pk_field.to_representation(doc_id)} 

def to_representation(self, value):
        if self.get_depth(self) == 0:
            return super(ComboReferenceField, self).to_representation(value)

        assert isinstance(value, (Document, DBRef))
        if isinstance(value, DBRef):
            value = self.model._get_db().dereference(value)

        ser = self.serializer(instance=value)
        return ser.data 

def to_representation(self, value):
        assert isinstance(value, (Document, DBRef))
        doc_id = value.id
        return self.pk_field.to_representation(doc_id) 

def mask_secrets(self, value):
        """
        Process the model dictionary and mask secret values.

        :type value: ``dict``
        :param value: Document dictionary.

        :rtype: ``dict``
        """
        return value 

def construct_fields(model, registry, only_fields, exclude_fields):
    """
    Args:
        model (mongoengine.Document):
        registry (graphene_mongo.registry.Registry):
        only_fields ([str]):
        exclude_fields ([str]):

    Returns:
        (OrderedDict, OrderedDict): converted fields and self reference fields.

    """
    _model_fields = get_model_fields(model)
    fields = OrderedDict()
    self_referenced = OrderedDict()
    for name, field in _model_fields.items():
        is_not_in_only = only_fields and name not in only_fields
        is_excluded = name in exclude_fields
        if is_not_in_only or is_excluded:
            # We skip this field if we specify only_fields and is not
            # in there. Or when we exclude this field in exclude_fields
            continue
        if isinstance(field, mongoengine.ListField):
            if not field.field:
                continue
            # Take care of list of self-reference.
            document_type_obj = field.field.__dict__.get("document_type_obj", None)
            if (
                document_type_obj == model._class_name
                or isinstance(document_type_obj, model)
                or document_type_obj == model
            ):
                self_referenced[name] = field
                continue
        converted = convert_mongoengine_field(field, registry)
        if not converted:
            continue
        fields[name] = converted

    return fields, self_referenced 

def is_valid_mongoengine_model(model):
    return inspect.isclass(model) and (
        issubclass(model, mongoengine.Document)
        or issubclass(model, mongoengine.EmbeddedDocument)
    ) 

def test_required_with_default(self):
        class Doc(me.Document):
            basic = me.IntField(required=True, default=42)
            cunning = me.BooleanField(required=True, default=False)
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
        doc, errors = DocSchema().load({})
        assert not errors
        assert doc.basic == 42
        assert doc.cunning is False 

def get_mongo_adaptor(original_class):
    if mongoengine is not None and issubclass(original_class, mongoengine.Document):
        return MongoAdaptor(original_class)
    return None 

def test_disable_skip_none_field(self):
        class Doc(me.Document):
            field_empty = me.StringField()
            list_empty = me.ListField(me.StringField())
        class DocSchema(ModelSchema):
            class Meta:
                model = Doc
                model_skip_values = ()
        doc = Doc()
        data, errors = DocSchema().dump(doc)
        assert not errors
        assert data == {'field_empty': None, 'list_empty': []} 

def test_kwargs(self):
        class MockModel(Document):
            foo = fields.StringField()
            bar = fields.StringField()
            baz = fields.StringField()

        class TestFS(ModelFilterset):
            class Meta:
                model = MockModel
                kwargs = {
                    'foo': { 'lookup': 'gte' }
                }
        fs = TestFS()
        self.assertEqual(fs.filters['foo'].lookup_type, 'gte') 

def test_auto_derivatives(self):
        class FooField(fields.StringField):
            pass

        class MockModel(Document):
            foo = FooField()

        class TestFS(ModelFilterset):
            class Meta:
                model = MockModel

        fs = TestFS()
        self.assertEqual(set(fs.filters.keys()), set(['id', 'foo']))
        self.assertIsInstance(fs.filters['foo'], filters.CharFilter) 

def test_custom_type(self):
        class FooField(fields.BaseField):
            pass

        class MockModel(Document):
            foo = FooField()

        class TestFS(ModelFilterset):
            filters_mapping = {
                FooField: filters.CharFilter
            }
            class Meta:
                model = MockModel
        fs = TestFS()
        self.assertEqual(set(fs.filters.keys()), set(['id', 'foo'])) 

def test_selecting(self):
        class MockModel(Document):
            foo = fields.StringField()
            bar = fields.StringField()
            baz = fields.StringField()

        class TestFS(ModelFilterset):
            class Meta:
                model = MockModel
                fields = ('foo','baz')

        fs = TestFS()
        self.assertEqual(set(fs.filters.keys()), set(['id', 'foo', 'baz']))