import pickle
import unittest

from mongoengine import Document, IntField, StringField
from mongoengine.connection import connect


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


class TestQuerysetPickable(unittest.TestCase):
    """
    Test for adding pickling support for QuerySet instances
    See issue https://github.com/MongoEngine/mongoengine/issues/442
    """

    def setUp(self):
        super(TestQuerysetPickable, self).setUp()
        connection = connect(db="test")
        connection.drop_database("test")
        self.john = Person.objects.create(name="John", age=21)

    def test_picke_simple_qs(self):

        qs = Person.objects.all()

        pickle.dumps(qs)

    def _get_loaded(self, qs):
        s = pickle.dumps(qs)

        return pickle.loads(s)

    def test_unpickle(self):
        qs = Person.objects.all()

        loadedQs = self._get_loaded(qs)

        self.assertEqual(qs.count(), loadedQs.count())

        # can update loadedQs
        loadedQs.update(age=23)

        # check
        self.assertEqual(Person.objects.first().age, 23)

    def test_pickle_support_filtration(self):
        Person.objects.create(name="Alice", age=22)

        Person.objects.create(name="Bob", age=23)

        qs = Person.objects.filter(age__gte=22)
        self.assertEqual(qs.count(), 2)

        loaded = self._get_loaded(qs)

        self.assertEqual(loaded.count(), 2)
        self.assertEqual(loaded.filter(name="Bob").first().age, 23)

@@ -123,16 +123,56 @@
 connect%0A
+from tests.utils import MongoDBTestCase%0A
 %0A%0Aclass 
@@ -268,17 +268,15 @@
 ble(
-unittest.
+MongoDB
 Test
@@ -502,89 +502,8 @@
 p()%0A
-        connection = connect(db=%22test%22)%0A        connection.drop_database(%22test%22)%0A
     

fix test_pickable that was brought back to life recently