Rollback Issue With TestCase

[seanpar203]

Following the documentation of decorating celery tasks with @close_connection, I've had some difficult with running django tests using pytest as the runner.

Here's the code(it's pretty simple)

test.py

    @patch("saleor.payment.models.Subscription.get_detail")
    def test_handle_invoice_payment_succeeded(self, retrieve_mock):
        retrieve_mock.return_value = Mock(
            status="active",
            current_period_end=datetime.now().timestamp(),
            current_period_start=datetime.now().timestamp(),
            cancel_at_period_end=False,
        )

        update_subscription(Events.invoice.payment_succeeded)

        subscription = Subscription.objects.get(
            stripe_id=Events.invoice.payment_succeeded
        )

        self.assertEqual(subscription.status, "active")

tasks.py

@app.task
@close_connection
def update_subscription(sub_id):
    """ Updates a subscription to match what's in Stripe. """
    try:
        subscription = Subscription.objects.get(stripe_id=sub_id)
    except Subscription.DoesNotExist:
        logger.exception("update_subscription_failure")
        return

    try:
        subscription.update()
    except StripeError:
        logger.exception("update_subscription_failure")
    else:
        subscription.save()

When running pytest I get these errors:

self = <django_db_geventpool.backends.postgresql_psycopg2.base.DatabaseWrapper object at 0x7f9f90bb1ac8>, rollback = True

    def set_rollback(self, rollback):
        """
        Set or unset the "needs rollback" flag -- for *advanced use* only.
        """
        if not self.in_atomic_block:
            raise TransactionManagementError(
>               "The rollback flag doesn't work outside of an 'atomic' block.")
E           django.db.transaction.TransactionManagementError: The rollback flag doesn't work outside of an 'atomic' block.

The test also fails but when I paste the task functionality inside the test and don't actually call the task function, the test passes.

Is this expected or am I doing something wrong?

[seanpar203]

I guess a solution is to create a function that does the work and have the task call that function, that way the function can be tested outside the scope of a task but still have the proper close_connection functionality when in production?

def calculate_money():
    # Do something
    return None

@app.task
@close_connection
def task_calculate_money():
    calculate_money()

Although this isn't ideal obviously.

[jneight]

Hi!, Have you tried with any of the others methods in the readme instead of using the decorator?

https://github.com/jneight/django-db-geventpool/blob/master/README.rst#using-orm-when-not-serving-requests

[michalc]

I had a similar issue when using pytest.mark.django_db. Ultimately, I didn't use close_connection that comes with django-db-geventpool, but rolled my own that only closed the connection if not in a transaction (which pytest.mark.django_db performs tests in). Something that did:

from django.db import connection

if not connection.in_atomic_block:
    connection.close_if_unusable_or_obsolete()

Not great, since I'm not such a fan having code specifically for the reason that it behaves differently in the tests than in production, but it may be the least-bad option.