# -*- coding: utf-8 -*- from mongoengine import * from tests.utils import MongoDBTestCase class TestSequenceField(MongoDBTestCase): def test_sequence_field(self): class Person(Document): id = SequenceField(primary_key=True) name = StringField() self.db["mongoengine.counters"].drop() Person.drop_collection() for x in range(10): Person(name="Person %s" % x).save() c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) ids = [i.id for i in Person.objects] self.assertEqual(ids, range(1, 11)) c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) Person.id.set_next_value(1000) c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 1000) def test_sequence_field_get_next_value(self): class Person(Document): id = SequenceField(primary_key=True) name = StringField() self.db["mongoengine.counters"].drop() Person.drop_collection() for x in range(10): Person(name="Person %s" % x).save() self.assertEqual(Person.id.get_next_value(), 11) self.db["mongoengine.counters"].drop() self.assertEqual(Person.id.get_next_value(), 1) class Person(Document): id = SequenceField(primary_key=True, value_decorator=str) name = StringField() self.db["mongoengine.counters"].drop() Person.drop_collection() for x in range(10): Person(name="Person %s" % x).save() self.assertEqual(Person.id.get_next_value(), "11") self.db["mongoengine.counters"].drop() self.assertEqual(Person.id.get_next_value(), "1") def test_sequence_field_sequence_name(self): class Person(Document): id = SequenceField(primary_key=True, sequence_name="jelly") name = StringField() self.db["mongoengine.counters"].drop() Person.drop_collection() for x in range(10): Person(name="Person %s" % x).save() c = self.db["mongoengine.counters"].find_one({"_id": "jelly.id"}) self.assertEqual(c["next"], 10) ids = [i.id for i in Person.objects] self.assertEqual(ids, range(1, 11)) c = self.db["mongoengine.counters"].find_one({"_id": "jelly.id"}) self.assertEqual(c["next"], 10) Person.id.set_next_value(1000) c = self.db["mongoengine.counters"].find_one({"_id": "jelly.id"}) self.assertEqual(c["next"], 1000) def test_multiple_sequence_fields(self): class Person(Document): id = SequenceField(primary_key=True) counter = SequenceField() name = StringField() self.db["mongoengine.counters"].drop() Person.drop_collection() for x in range(10): Person(name="Person %s" % x).save() c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) ids = [i.id for i in Person.objects] self.assertEqual(ids, range(1, 11)) counters = [i.counter for i in Person.objects] self.assertEqual(counters, range(1, 11)) c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) Person.id.set_next_value(1000) c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 1000) Person.counter.set_next_value(999) c = self.db["mongoengine.counters"].find_one({"_id": "person.counter"}) self.assertEqual(c["next"], 999) def test_sequence_fields_reload(self): class Animal(Document): counter = SequenceField() name = StringField() self.db["mongoengine.counters"].drop() Animal.drop_collection() a = Animal(name="Boi").save() self.assertEqual(a.counter, 1) a.reload() self.assertEqual(a.counter, 1) a.counter = None self.assertEqual(a.counter, 2) a.save() self.assertEqual(a.counter, 2) a = Animal.objects.first() self.assertEqual(a.counter, 2) a.reload() self.assertEqual(a.counter, 2) def test_multiple_sequence_fields_on_docs(self): class Animal(Document): id = SequenceField(primary_key=True) name = StringField() class Person(Document): id = SequenceField(primary_key=True) name = StringField() self.db["mongoengine.counters"].drop() Animal.drop_collection() Person.drop_collection() for x in range(10): Animal(name="Animal %s" % x).save() Person(name="Person %s" % x).save() c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) c = self.db["mongoengine.counters"].find_one({"_id": "animal.id"}) self.assertEqual(c["next"], 10) ids = [i.id for i in Person.objects] self.assertEqual(ids, range(1, 11)) id = [i.id for i in Animal.objects] self.assertEqual(id, range(1, 11)) c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) c = self.db["mongoengine.counters"].find_one({"_id": "animal.id"}) self.assertEqual(c["next"], 10) def test_sequence_field_value_decorator(self): class Person(Document): id = SequenceField(primary_key=True, value_decorator=str) name = StringField() self.db["mongoengine.counters"].drop() Person.drop_collection() for x in range(10): p = Person(name="Person %s" % x) p.save() c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) ids = [i.id for i in Person.objects] self.assertEqual(ids, map(str, range(1, 11))) c = self.db["mongoengine.counters"].find_one({"_id": "person.id"}) self.assertEqual(c["next"], 10) def test_embedded_sequence_field(self): class Comment(EmbeddedDocument): id = SequenceField() content = StringField(required=True) class Post(Document): title = StringField(required=True) comments = ListField(EmbeddedDocumentField(Comment)) self.db["mongoengine.counters"].drop() Post.drop_collection() Post( title="MongoEngine", comments=[ Comment(content="NoSQL Rocks"), Comment(content="MongoEngine Rocks"), ], ).save() c = self.db["mongoengine.counters"].find_one({"_id": "comment.id"}) self.assertEqual(c["next"], 2) post = Post.objects.first() self.assertEqual(1, post.comments[0].id) self.assertEqual(2, post.comments[1].id) def test_inherited_sequencefield(self): class Base(Document): name = StringField() counter = SequenceField() meta = {"abstract": True} class Foo(Base): pass class Bar(Base): pass bar = Bar(name="Bar") bar.save() foo = Foo(name="Foo") foo.save() self.assertTrue( "base.counter" in self.db["mongoengine.counters"].find().distinct("_id") ) self.assertFalse( ("foo.counter" or "bar.counter") in self.db["mongoengine.counters"].find().distinct("_id") ) self.assertNotEqual(foo.counter, bar.counter) self.assertEqual(foo._fields["counter"].owner_document, Base) self.assertEqual(bar._fields["counter"].owner_document, Base) def test_no_inherited_sequencefield(self): class Base(Document): name = StringField() meta = {"abstract": True} class Foo(Base): counter = SequenceField() class Bar(Base): counter = SequenceField() bar = Bar(name="Bar") bar.save() foo = Foo(name="Foo") foo.save() self.assertFalse( "base.counter" in self.db["mongoengine.counters"].find().distinct("_id") ) self.assertTrue( ("foo.counter" and "bar.counter") in self.db["mongoengine.counters"].find().distinct("_id") ) self.assertEqual(foo.counter, bar.counter) self.assertEqual(foo._fields["counter"].owner_document, Foo) self.assertEqual(bar._fields["counter"].owner_document, Bar)