<filename>tenant_schemas_celery/tests.py
from django.db import connection
from django.utils.unittest import skipIf


from tenant_schemas.tests.models import Tenant, DummyModel
from tenant_schemas.tests.testcases import BaseTestCase
from tenant_schemas.utils import get_public_schema_name

try:
    from .app import CeleryApp
except ImportError:
    app = None
else:
    app = CeleryApp('testapp')

    class CeleryConfig:
        CELERY_ALWAYS_EAGER = True
        CELERY_EAGER_PROPAGATES_EXCEPTIONS = True

    app.config_from_object(CeleryConfig)

    @app.task
    def update_task(model_id, name):
        dummy = DummyModel.objects.get(pk=model_id)
        dummy.name = name
        dummy.save()

    @app.task
    def update_retry_task(model_id, name):
        if update_retry_task.request.retries:
            return update_task(model_id, name)

        # Don't throw the Retry exception.
        update_retry_task.retry(throw=False)


@skipIf(app is None, 'Celery is not available.')
class CeleryTasksTests(BaseTestCase):
    def setUp(self):
        super(CeleryTasksTests, self).setUp()
        self.tenant1 = Tenant(domain_url='test1', schema_name='test1')
        self.tenant1.save()

        self.tenant2 = Tenant(domain_url='test2', schema_name='test2')
        self.tenant2.save()

        connection.set_tenant(self.tenant1)
        self.dummy1 = DummyModel.objects.create(name='test1')

        connection.set_tenant(self.tenant2)
        self.dummy2 = DummyModel.objects.create(name='test2')

        connection.set_schema_to_public()

    def test_basic_model_update(self):
        # We should be in public schema where dummies don't exist.
        for dummy in self.dummy1, self.dummy2:
            # Test both async and local versions.
            with self.assertRaises(DummyModel.DoesNotExist):
                update_task.apply_async(args=(dummy.pk, 'updated-name'))

            with self.assertRaises(DummyModel.DoesNotExist):
                update_task.apply(args=(dummy.pk, 'updated-name'))

        connection.set_tenant(self.tenant1)
        update_task.apply_async(args=(self.dummy1.pk, 'updated-name'))
        self.assertEqual(connection.schema_name, self.tenant1.schema_name)

        # The task restores the schema from before running the task, so we are
        # using the `tenant1` tenant now.
        model_count = DummyModel.objects.filter(name='updated-name').count()
        self.assertEqual(model_count, 1)

        connection.set_tenant(self.tenant2)
        model_count = DummyModel.objects.filter(name='updated-name').count()
        self.assertEqual(model_count, 0)

    def test_task_retry(self):
        # Schema name should persist through retry attempts.
        connection.set_tenant(self.tenant1)
        update_retry_task.apply_async(args=(self.dummy1.pk, 'updated-name'))

        model_count = DummyModel.objects.filter(name='updated-name').count()
        self.assertEqual(model_count, 1)

    def test_restoring_schema_name(self):
        update_task.apply_async(
            args=(self.dummy1.pk, 'updated-name'),
            kwargs={'_schema_name': self.tenant1.schema_name}
        )
        self.assertEqual(connection.schema_name, get_public_schema_name())

        connection.set_tenant(self.tenant1)
        update_task.apply_async(
            args=(self.dummy2.pk, 'updated-name'),
            kwargs={'_schema_name': self.tenant2.schema_name}
        )
        self.assertEqual(connection.schema_name, self.tenant1.schema_name)

        connection.set_tenant(self.tenant2)
        # The model does not exist in the public schema.
        with self.assertRaises(DummyModel.DoesNotExist):
            update_task.apply_async(
                args=(self.dummy2.pk, 'updated-name'),
                kwargs={'_schema_name': get_public_schema_name()}
            )

        self.assertEqual(connection.schema_name, self.tenant2.schema_name)

<filename>tenant_schemas_celery/tests.py
from django.db import connection
from django.utils.unittest import skipIf


from tenant_schemas.tests.models import Tenant, DummyModel
from tenant_schemas.tests.testcases import BaseTestCase
from tenant_schemas.utils import get_public_schema_name

try:
    from .app import CeleryApp
except ImportError:
    app = None
else:
    app = CeleryApp('testapp')

    class CeleryConfig:
        CELERY_ALWAYS_EAGER = True
        CELERY_EAGER_PROPAGATES_EXCEPTIONS = True

    app.config_from_object(CeleryConfig)

    @app.task
    def update_task(model_id, name):
        dummy = DummyModel.objects.get(pk=model_id)
        dummy.name = name
        dummy.save()

    @app.task
    def update_retry_task(model_id, name):
        if update_retry_task.request.retries:
            return update_task(model_id, name)

        # Don't throw the Retry exception.
        update_retry_task.retry(throw=False)


@skipIf(app is None, 'Celery is not available.')
class CeleryTasksTests(BaseTestCase):
    def setUp(self):
        super(CeleryTasksTests, self).setUp()
        self.tenant1 = Tenant(domain_url='test1', schema_name='test1')
        self.tenant1.save()

        self.tenant2 = Tenant(domain_url='test2', schema_name='test2')
        self.tenant2.save()

        connection.set_tenant(self.tenant1)
        self.dummy1 = DummyModel.objects.create(name='test1')

        connection.set_tenant(self.tenant2)
        self.dummy2 = DummyModel.objects.create(name='test2')

        connection.set_schema_to_public()

    def test_basic_model_update(self):
        # We should be in public schema where dummies don't exist.
        for dummy in self.dummy1, self.dummy2:
            # Test both async and local versions.
            with self.assertRaises(DummyModel.DoesNotExist):
                update_task.apply_async(args=(dummy.pk, 'updated-name'))

            with self.assertRaises(DummyModel.DoesNotExist):
                update_task.apply(args=(dummy.pk, 'updated-name'))

        connection.set_tenant(self.tenant1)
        update_task.apply_async(args=(self.dummy1.pk, 'updated-name'))
        self.assertEqual(connection.schema_name, self.tenant1.schema_name)

        # The task restores the schema from before running the task, so we are
        # using the `tenant1` tenant now.
        model_count = DummyModel.objects.filter(name='updated-name').count()
        self.assertEqual(model_count, 1)

        connection.set_tenant(self.tenant2)
        model_count = DummyModel.objects.filter(name='updated-name').count()
        self.assertEqual(model_count, 0)

    def test_task_retry(self):
        # Schema name should persist through retry attempts.
        connection.set_tenant(self.tenant1)
        update_retry_task.apply_async(args=(self.dummy1.pk, 'updated-name'))

        model_count = DummyModel.objects.filter(name='updated-name').count()
        self.assertEqual(model_count, 1)

    def test_restoring_schema_name(self):
        update_task.apply_async(
            args=(self.dummy1.pk, 'updated-name'),
            kwargs={'_schema_name': self.tenant1.schema_name}
        )
        self.assertEqual(connection.schema_name, get_public_schema_name())

        connection.set_tenant(self.tenant1)
        update_task.apply_async(
            args=(self.dummy2.pk, 'updated-name'),
            kwargs={'_schema_name': self.tenant2.schema_name}
        )
        self.assertEqual(connection.schema_name, self.tenant1.schema_name)

        connection.set_tenant(self.tenant2)
        # The model does not exist in the public schema.
        with self.assertRaises(DummyModel.DoesNotExist):
            update_task.apply_async(
                args=(self.dummy2.pk, 'updated-name'),
                kwargs={'_schema_name': get_public_schema_name()}
            )

        self.assertEqual(connection.schema_name, self.tenant2.schema_name)

# Don't throw the Retry exception. # We should be in public schema where dummies don't exist. # Test both async and local versions. # The task restores the schema from before running the task, so we are # using the `tenant1` tenant now. # Schema name should persist through retry attempts. # The model does not exist in the public schema.