1# Copyright: See the LICENSE file.
2
3
4"""factory_boy extensions for use with the mongoengine library (pymongo wrapper)."""
5
6
7from . import base
8
9
10class MongoEngineFactory(base.Factory):
11    """Factory for mongoengine objects."""
12
13    class Meta:
14        abstract = True
15
16    @classmethod
17    def _build(cls, model_class, *args, **kwargs):
18        return model_class(*args, **kwargs)
19
20    @classmethod
21    def _create(cls, model_class, *args, **kwargs):
22        instance = model_class(*args, **kwargs)
23        if instance._is_document:
24            instance.save()
25        return instance
26