Source

djangoappengine / tests / batch.py

from __future__ import with_statement

from django.dispatch import receiver
from django.db import models
from django.db.models.signals import pre_save
from django.db.models.signals import post_save
from django.db.models.signals import pre_delete
from django.db.models.signals import post_delete
from django.db.models import BatchOperation
from django.test import TestCase

from google.appengine.api import datastore
from google.appengine.api.datastore import Put
from google.appengine.api.datastore import Delete

from .testmodels import Parent
from .testmodels import Mti


class mem(object):
    """Global AppEngine datastore operations counters"""
    put_count = 0
    delete_count = 0


def new_put(entities, **kwargs):
    mem.put_count += 1
    return Put(entities, **kwargs)


def new_delete(keys, **kwargs):
    mem.delete_count += 1
    return Delete(keys, **kwargs)


class BatchCategory(models.Model):
    name = models.CharField(max_length=30)


@receiver(pre_save, sender=BatchCategory)
def batch_category_pre_save(sender, instance=None, **kwargs):
    instance.batch_category_pre_save = True


@receiver(post_save, sender=BatchCategory)
def batch_category_post_save(sender, instance=None, **kwargs):
    instance.batch_category_post_save = True


@receiver(pre_delete, sender=BatchCategory)
def batch_category_pre_delete(sender, instance=None, **kwargs):
    instance.batch_category_pre_delete = True


@receiver(post_delete, sender=BatchCategory)
def batch_category_post_delete(sender, instance=None, **kwargs):
    instance.batch_category_post_delete = True


class BatchTest(TestCase):
    def setUp(self):
        # reset AppEngine datastore operations counters
        mem.put_count = 0
        mem.delete_count = 0

        datastore.Put = new_put
        datastore.Delete = new_delete

    def tearDown(self):
        datastore.Put = Put
        datastore.Delete = Delete

    def test_batch_save(self):
        with BatchOperation() as q:
            categories = []
            for name in ['appengine', 'django']:
                c = BatchCategory(name=name)
                q.save(c)
                categories.append(c)

                self.assertTrue(c.batch_category_pre_save)
                self.assertEqual(BatchCategory.objects.count(), 0)

        for c in categories:
            self.assertTrue(c.batch_category_post_save)

        self.assertEqual(BatchCategory.objects.count(), 2)
        self.assertEqual(mem.put_count, 1)

        with BatchOperation() as q:
            categories = BatchCategory.objects.all()
            for c in categories:
                q.delete(c)

                self.assertTrue(c.batch_category_pre_delete)
                self.assertEqual(BatchCategory.objects.count(), 2)

        for c in categories:
            self.assertTrue(c.batch_category_post_delete)

        self.assertEqual(BatchCategory.objects.count(), 0)
        self.assertEqual(mem.delete_count, 1)

    def test_plain_operation_in_batch(self):
        with BatchOperation() as q:
            appengine = BatchCategory(name='appengine')
            q.save(appengine)
            self.assertEqual(BatchCategory.objects.count(), 0)

            django = BatchCategory(name='django')
            django.save()
            self.assertEqual(BatchCategory.objects.count(), 1)

        self.assertIsInstance(appengine.pk, int)
        self.assertIsInstance(appengine.id, int)
        self.assertEqual(appengine.pk, appengine.id)

        self.assertIsInstance(django.pk, int)
        self.assertIsInstance(django.id, int)
        self.assertEqual(django.pk, django.id)

        self.assertEqual(BatchCategory.objects.count(), 2)
        self.assertEqual(mem.put_count, 2)

        appengine = BatchCategory.objects.get(name='appengine')
        django = BatchCategory.objects.get(name='django')
        with BatchOperation() as q:
            q.delete(appengine)
            self.assertEqual(BatchCategory.objects.count(), 2)

            django.delete()
            self.assertEqual(BatchCategory.objects.count(), 1)

        self.assertEqual(appengine.pk, None)
        self.assertEqual(django.pk, None)

        self.assertEqual(BatchCategory.objects.count(), 0)
        self.assertEqual(mem.delete_count, 2)

    def test_batch_operation_params(self):
        """Test, that pool is flushed according to config"""
        c1 = BatchCategory(name='test1')
        c2 = BatchCategory(name='test2')

        config = {'default': {'pool_size': 2,
                              'batch_size': 2,
                              'datastore_deadline': 30}}
        with BatchOperation(config) as op:
            op.save(c1)
            self.assertEqual(c1.id, None, 'Pool is already flushed')

            op.save(c2)
            self.assertNotEqual(c1.id, None, 'Pool is not flushed yet')
            self.assertNotEqual(c2.id, None, 'Pool is not flushed yet')

        self.assertEqual(BatchCategory.objects.count(), 2)


class BatchMtiTest(TestCase):
    def test_mti(self):
        with BatchOperation() as op:
            op.save(Mti(foo=4, bar=2))
            op.save(Mti(foo=5, bar=3))

        self.assertEqual(Parent.objects.count(), 2)
        self.assertEqual(Mti.objects.only('bar').count(), 2)

        for parent in Parent.objects.all():
            print 'parent.id', parent.id

        for mti in Mti.objects.only('id', 'bar'):
            parent = Parent.objects.get(id=mti.id)

            print mti.id
            print mti.pk
            print mti.parent_ptr_id
            print mti._meta.pk.attname
            #parent = mti.parent_ptr
            #self.assertEqual(mti.bar, parent.foo - 2)
        self.fail()

    def test_mti_with_params(self):
        """Test, that pool is flushed according to config"""
        mti1 = Mti(foo=2, bar=0)
        mti2 = Mti(foo=4, bar=2)

        config = {'default': {'pool_size': 4,
                              'batch_size': 8,
                              'datastore_deadline': 30}}
        with BatchOperation(config) as op:
            op.save(mti1)
            self.assertEqual(mti1.id, None, 'Pool is already flushed')

            op.save(mti2)
            self.assertNotEqual(mti1.id, None, 'Pool is not flushed yet')
            self.assertNotEqual(mti2.id, None, 'Pool is not flushed yet')