# -*- coding: utf-8 -*-
import unittest

from mongoengine import *
from mongoengine.pymongo_support import list_collection_names

from mongoengine.queryset import NULLIFY, PULL
from mongoengine.connection import get_db

__all__ = ("ClassMethodsTest",)


class ClassMethodsTest(unittest.TestCase):
    def setUp(self):
        connect(db="mongoenginetest")
        self.db = get_db()

        class Person(Document):
            name = StringField()
            age = IntField()

            non_field = True

            meta = {"allow_inheritance": True}

        self.Person = Person

    def tearDown(self):
        for collection in list_collection_names(self.db):
            self.db.drop_collection(collection)

    def test_definition(self):
        """Ensure that document may be defined using fields.
        """
        self.assertEqual(
            ["_cls", "age", "id", "name"], sorted(self.Person._fields.keys())
        )
        self.assertEqual(
            ["IntField", "ObjectIdField", "StringField", "StringField"],
            sorted([x.__class__.__name__ for x in self.Person._fields.values()]),
        )

    def test_get_db(self):
        """Ensure that get_db returns the expected db.
        """
        db = self.Person._get_db()
        self.assertEqual(self.db, db)

    def test_get_collection_name(self):
        """Ensure that get_collection_name returns the expected collection
        name.
        """
        collection_name = "person"
        self.assertEqual(collection_name, self.Person._get_collection_name())

    def test_get_collection(self):
        """Ensure that get_collection returns the expected collection.
        """
        collection_name = "person"
        collection = self.Person._get_collection()
        self.assertEqual(self.db[collection_name], collection)

    def test_drop_collection(self):
        """Ensure that the collection may be dropped from the database.
        """
        collection_name = "person"
        self.Person(name="Test").save()
        self.assertIn(collection_name, list_collection_names(self.db))

        self.Person.drop_collection()
        self.assertNotIn(collection_name, list_collection_names(self.db))

    def test_register_delete_rule(self):
        """Ensure that register delete rule adds a delete rule to the document
        meta.
        """

        class Job(Document):
            employee = ReferenceField(self.Person)

        self.assertEqual(self.Person._meta.get("delete_rules"), None)

        self.Person.register_delete_rule(Job, "employee", NULLIFY)
        self.assertEqual(
            self.Person._meta["delete_rules"], {(Job, "employee"): NULLIFY}
        )

    def test_compare_indexes(self):
        """ Ensure that the indexes are properly created and that
        compare_indexes identifies the missing/extra indexes
        """

        class BlogPost(Document):
            author = StringField()
            title = StringField()
            description = StringField()
            tags = StringField()

            meta = {"indexes": [("author", "title")]}

        BlogPost.drop_collection()

        BlogPost.ensure_indexes()
        self.assertEqual(BlogPost.compare_indexes(), {"missing": [], "extra": []})

        BlogPost.ensure_index(["author", "description"])
        self.assertEqual(
            BlogPost.compare_indexes(),
            {"missing": [], "extra": [[("author", 1), ("description", 1)]]},
        )

        BlogPost._get_collection().drop_index("author_1_description_1")
        self.assertEqual(BlogPost.compare_indexes(), {"missing": [], "extra": []})

        BlogPost._get_collection().drop_index("author_1_title_1")
        self.assertEqual(
            BlogPost.compare_indexes(),
            {"missing": [[("author", 1), ("title", 1)]], "extra": []},
        )

    def test_compare_indexes_inheritance(self):
        """ Ensure that the indexes are properly created and that
        compare_indexes identifies the missing/extra indexes for subclassed
        documents (_cls included)
        """

        class BlogPost(Document):
            author = StringField()
            title = StringField()
            description = StringField()

            meta = {"allow_inheritance": True}

        class BlogPostWithTags(BlogPost):
            tags = StringField()
            tag_list = ListField(StringField())

            meta = {"indexes": [("author", "tags")]}

        BlogPost.drop_collection()

        BlogPost.ensure_indexes()
        BlogPostWithTags.ensure_indexes()
        self.assertEqual(BlogPost.compare_indexes(), {"missing": [], "extra": []})

        BlogPostWithTags.ensure_index(["author", "tag_list"])
        self.assertEqual(
            BlogPost.compare_indexes(),
            {"missing": [], "extra": [[("_cls", 1), ("author", 1), ("tag_list", 1)]]},
        )

        BlogPostWithTags._get_collection().drop_index("_cls_1_author_1_tag_list_1")
        self.assertEqual(BlogPost.compare_indexes(), {"missing": [], "extra": []})

        BlogPostWithTags._get_collection().drop_index("_cls_1_author_1_tags_1")
        self.assertEqual(
            BlogPost.compare_indexes(),
            {"missing": [[("_cls", 1), ("author", 1), ("tags", 1)]], "extra": []},
        )

    def test_compare_indexes_multiple_subclasses(self):
        """ Ensure that compare_indexes behaves correctly if called from a
        class, which base class has multiple subclasses
        """

        class BlogPost(Document):
            author = StringField()
            title = StringField()
            description = StringField()

            meta = {"allow_inheritance": True}

        class BlogPostWithTags(BlogPost):
            tags = StringField()
            tag_list = ListField(StringField())

            meta = {"indexes": [("author", "tags")]}

        class BlogPostWithCustomField(BlogPost):
            custom = DictField()

            meta = {"indexes": [("author", "custom")]}

        BlogPost.ensure_indexes()
        BlogPostWithTags.ensure_indexes()
        BlogPostWithCustomField.ensure_indexes()

        self.assertEqual(BlogPost.compare_indexes(), {"missing": [], "extra": []})
        self.assertEqual(
            BlogPostWithTags.compare_indexes(), {"missing": [], "extra": []}
        )
        self.assertEqual(
            BlogPostWithCustomField.compare_indexes(), {"missing": [], "extra": []}
        )

    def test_compare_indexes_for_text_indexes(self):
        """ Ensure that compare_indexes behaves correctly for text indexes """

        class Doc(Document):
            a = StringField()
            b = StringField()
            meta = {
                "indexes": [
                    {
                        "fields": ["$a", "$b"],
                        "default_language": "english",
                        "weights": {"a": 10, "b": 2},
                    }
                ]
            }

        Doc.drop_collection()
        Doc.ensure_indexes()
        actual = Doc.compare_indexes()
        expected = {"missing": [], "extra": []}
        self.assertEqual(actual, expected)

    def test_list_indexes_inheritance(self):
        """ ensure that all of the indexes are listed regardless of the super-
        or sub-class that we call it from
        """

        class BlogPost(Document):
            author = StringField()
            title = StringField()
            description = StringField()

            meta = {"allow_inheritance": True}

        class BlogPostWithTags(BlogPost):
            tags = StringField()

            meta = {"indexes": [("author", "tags")]}

        class BlogPostWithTagsAndExtraText(BlogPostWithTags):
            extra_text = StringField()

            meta = {"indexes": [("author", "tags", "extra_text")]}

        BlogPost.drop_collection()

        BlogPost.ensure_indexes()
        BlogPostWithTags.ensure_indexes()
        BlogPostWithTagsAndExtraText.ensure_indexes()

        self.assertEqual(BlogPost.list_indexes(), BlogPostWithTags.list_indexes())
        self.assertEqual(
            BlogPost.list_indexes(), BlogPostWithTagsAndExtraText.list_indexes()
        )
        self.assertEqual(
            BlogPost.list_indexes(),
            [
                [("_cls", 1), ("author", 1), ("tags", 1)],
                [("_cls", 1), ("author", 1), ("tags", 1), ("extra_text", 1)],
                [(u"_id", 1)],
                [("_cls", 1)],
            ],
        )

    def test_register_delete_rule_inherited(self):
        class Vaccine(Document):
            name = StringField(required=True)

            meta = {"indexes": ["name"]}

        class Animal(Document):
            family = StringField(required=True)
            vaccine_made = ListField(
                ReferenceField("Vaccine", reverse_delete_rule=PULL)
            )

            meta = {"allow_inheritance": True, "indexes": ["family"]}

        class Cat(Animal):
            name = StringField(required=True)

        self.assertEqual(Vaccine._meta["delete_rules"][(Animal, "vaccine_made")], PULL)
        self.assertEqual(Vaccine._meta["delete_rules"][(Cat, "vaccine_made")], PULL)

    def test_collection_naming(self):
        """Ensure that a collection with a specified name may be used.
        """

        class DefaultNamingTest(Document):
            pass

        self.assertEqual(
            "default_naming_test", DefaultNamingTest._get_collection_name()
        )

        class CustomNamingTest(Document):
            meta = {"collection": "pimp_my_collection"}

        self.assertEqual("pimp_my_collection", CustomNamingTest._get_collection_name())

        class DynamicNamingTest(Document):
            meta = {"collection": lambda c: "DYNAMO"}

        self.assertEqual("DYNAMO", DynamicNamingTest._get_collection_name())

        # Use Abstract class to handle backwards compatibility
        class BaseDocument(Document):
            meta = {"abstract": True, "collection": lambda c: c.__name__.lower()}

        class OldNamingConvention(BaseDocument):
            pass

        self.assertEqual(
            "oldnamingconvention", OldNamingConvention._get_collection_name()
        )

        class InheritedAbstractNamingTest(BaseDocument):
            meta = {"collection": "wibble"}

        self.assertEqual("wibble", InheritedAbstractNamingTest._get_collection_name())

        # Mixin tests
        class BaseMixin(object):
            meta = {"collection": lambda c: c.__name__.lower()}

        class OldMixinNamingConvention(Document, BaseMixin):
            pass

        self.assertEqual(
            "oldmixinnamingconvention", OldMixinNamingConvention._get_collection_name()
        )

        class BaseMixin(object):
            meta = {"collection": lambda c: c.__name__.lower()}

        class BaseDocument(Document, BaseMixin):
            meta = {"allow_inheritance": True}

        class MyDocument(BaseDocument):
            pass

        self.assertEqual("basedocument", MyDocument._get_collection_name())

    def test_custom_collection_name_operations(self):
        """Ensure that a collection with a specified name is used as expected.
        """
        collection_name = "personCollTest"

        class Person(Document):
            name = StringField()
            meta = {"collection": collection_name}

        Person(name="Test User").save()
        self.assertIn(collection_name, list_collection_names(self.db))

        user_obj = self.db[collection_name].find_one()
        self.assertEqual(user_obj["name"], "Test User")

        user_obj = Person.objects[0]
        self.assertEqual(user_obj.name, "Test User")

        Person.drop_collection()
        self.assertNotIn(collection_name, list_collection_names(self.db))

    def test_collection_name_and_primary(self):
        """Ensure that a collection with a specified name may be used.
        """

        class Person(Document):
            name = StringField(primary_key=True)
            meta = {"collection": "app"}

        Person(name="Test User").save()

        user_obj = Person.objects.first()
        self.assertEqual(user_obj.name, "Test User")

        Person.drop_collection()


if __name__ == "__main__":
    unittest.main()