import datetime
import json
import pickle
import re
from unittest import mock

import celery
import pytest
from celery import states, uuid
from celery.app.task import Context
from celery.result import AsyncResult, GroupResult
from celery.utils.serialization import b64decode
from celery.worker.request import Request
from celery.worker.strategy import hybrid_to_proto2
from django.test import TransactionTestCase

from django_celery_results.backends.database import DatabaseBackend
from django_celery_results.models import ChordCounter, TaskResult


class SomeClass:

    def __init__(self, data):
        self.data = data


@pytest.mark.django_db()
@pytest.mark.usefixtures('depends_on_current_app')
class test_DatabaseBackend:

    @pytest.fixture(autouse=True)
    def setup_backend(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.result_backend = (
            'django_celery_results.backends:DatabaseBackend')
        self.app.conf.result_extended = True
        self.b = DatabaseBackend(app=self.app)

    def _create_request(self, task_id, name, args, kwargs,
                        argsrepr=None, kwargsrepr=None, task_protocol=2):
        msg = self.app.amqp.task_protocols[task_protocol](
            task_id=task_id,
            name=name,
            args=args,
            kwargs=kwargs,
            argsrepr=argsrepr,
            kwargsrepr=kwargsrepr,
        )
        if task_protocol == 1:
            body, headers, _, _ = hybrid_to_proto2(msg, msg.body)
            properties = None
            sent_event = {}
        else:
            headers, properties, body, sent_event = msg
        context = Context(
            headers=headers,
            properties=properties,
            body=body,
            sent_event=sent_event,
        )
        request = Request(context, decoded=True, task=name)
        if task_protocol == 1:
            assert request.argsrepr is None
            assert request.kwargsrepr is None
        else:
            assert request.argsrepr is not None
            assert request.kwargsrepr is not None
        return request

    def test_backend__pickle_serialization__dict_result(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, SomeClass(67)],
            kwargs={'c': 6, 'd': 'e', 'f': SomeClass(89)},
        )
        result = {'foo': 'baz', 'bar': SomeClass(12345)}

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result').get('foo') == 'baz'
        assert mindb.get('result').get('bar').data == 12345
        assert len(mindb.get('worker')) > 1
        assert mindb.get('task_name') == 'my_task'
        assert bool(re.match(
            r"\['a', 1, <.*SomeClass object at .*>\]",
            mindb.get('task_args')
        ))
        assert bool(re.match(
            r"{'c': 6, 'd': 'e', 'f': <.*SomeClass object at .*>}",
            mindb.get('task_kwargs')
        ))

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        task_args = pickle.loads(b64decode(tr.task_args))
        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_args == mindb.get('task_args')
        assert task_kwargs == mindb.get('task_kwargs')

        # check async_result
        ar = AsyncResult(tid2)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

        # check backward compatibility
        task_kwargs2 = str(request.kwargs)
        task_args2 = str(request.args)
        assert tr.task_args != task_args2
        assert tr.task_kwargs != task_kwargs2
        tr.task_args = task_args2
        tr.task_kwargs = task_kwargs2
        tr.save()
        mindb = self.b.get_task_meta(tid2)
        assert bool(re.match(
            r"\['a', 1, <.*SomeClass object at .*>\]",
            mindb.get('task_args')
        ))
        assert bool(re.match(
            r"{'c': 6, 'd': 'e', 'f': <.*SomeClass object at .*>}",
            mindb.get('task_kwargs')
        ))
        ar = AsyncResult(tid2)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

        tid3 = uuid()
        try:
            raise KeyError('foo')
        except KeyError as exception:
            self.b.mark_as_failure(tid3, exception)

        assert self.b.get_status(tid3) == states.FAILURE
        assert isinstance(self.b.get_result(tid3), KeyError)

    def test_backend__pickle_serialization__str_result(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, SomeClass(67)],
            kwargs={'c': 6, 'd': 'e', 'f': SomeClass(89)},
        )
        result = 'foo'

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result') == 'foo'
        assert mindb.get('task_name') == 'my_task'
        assert len(mindb.get('worker')) > 1
        assert bool(re.match(
            r"\['a', 1, <.*SomeClass object at .*>\]",
            mindb.get('task_args')
        ))
        assert bool(re.match(
            r"{'c': 6, 'd': 'e', 'f': <.*SomeClass object at .*>}",
            mindb.get('task_kwargs')
        ))

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        task_args = pickle.loads(b64decode(tr.task_args))
        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_args == mindb.get('task_args')
        assert task_kwargs == mindb.get('task_kwargs')

        # check async_result
        ar = AsyncResult(tid2)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

    def test_backend__pickle_serialization__bytes_result(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, SomeClass(67)],
            kwargs={'c': 6, 'd': 'e', 'f': SomeClass(89)},
        )
        result = b'foo'

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result') == b'foo'
        assert mindb.get('task_name') == 'my_task'
        assert len(mindb.get('worker')) > 1
        assert bool(re.match(
            r"\['a', 1, <.*SomeClass object at .*>\]",
            mindb.get('task_args')
        ))
        assert bool(re.match(
            r"{'c': 6, 'd': 'e', 'f': <.*SomeClass object at .*>}",
            mindb.get('task_kwargs')
        ))

        # check task_result objects
        tr = TaskResult.objects.get(task_id=tid2)
        task_args = pickle.loads(b64decode(tr.task_args))
        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_args == mindb.get('task_args')
        assert task_kwargs == mindb.get('task_kwargs')

        # check async_result
        ar = AsyncResult(tid2)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

    def test_backend__json_serialization__dict_result(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, True],
            kwargs={'c': 6, 'd': 'e', 'f': False},
        )
        result = {'foo': 'baz', 'bar': True}

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result').get('foo') == 'baz'
        assert mindb.get('result').get('bar') is True
        assert mindb.get('task_name') == 'my_task'
        assert mindb.get('task_args') == "['a', 1, True]"
        assert mindb.get('task_kwargs') == "{'c': 6, 'd': 'e', 'f': False}"

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        assert json.loads(tr.task_args) == "['a', 1, True]"
        assert json.loads(tr.task_kwargs) == "{'c': 6, 'd': 'e', 'f': False}"

        # check async_result
        ar = AsyncResult(tid2)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

        # check backward compatibility
        task_kwargs2 = str(request.kwargs)
        task_args2 = str(request.args)
        assert tr.task_args != task_args2
        assert tr.task_kwargs != task_kwargs2
        tr.task_args = task_args2
        tr.task_kwargs = task_kwargs2
        tr.save()
        mindb = self.b.get_task_meta(tid2)
        assert mindb.get('task_args') == "['a', 1, True]"
        assert mindb.get('task_kwargs') == "{'c': 6, 'd': 'e', 'f': False}"
        ar = AsyncResult(tid2)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

        tid3 = uuid()
        try:
            raise KeyError('foo')
        except KeyError as exception:
            self.b.mark_as_failure(tid3, exception)

        assert self.b.get_status(tid3) == states.FAILURE
        assert isinstance(self.b.get_result(tid3), KeyError)

    def test_backend__json_serialization__str_result(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, True],
            kwargs={'c': 6, 'd': 'e', 'f': False},
        )
        result = 'foo'

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result') == 'foo'
        assert mindb.get('task_name') == 'my_task'
        assert mindb.get('task_args') == "['a', 1, True]"
        assert mindb.get('task_kwargs') == "{'c': 6, 'd': 'e', 'f': False}"

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        assert json.loads(tr.task_args) == "['a', 1, True]"
        assert json.loads(tr.task_kwargs) == "{'c': 6, 'd': 'e', 'f': False}"

        # check async_result
        ar = AsyncResult(tid2)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

    def test_backend__pickle_serialization__dict_result__protocol_1(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, SomeClass(67)],
            kwargs={'c': 6, 'd': 'e', 'f': SomeClass(89)},
            task_protocol=1,
        )
        result = {'foo': 'baz', 'bar': SomeClass(12345)}

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result').get('foo') == 'baz'
        assert mindb.get('result').get('bar').data == 12345
        assert mindb.get('task_name') == 'my_task'

        assert mindb.get('task_args')[0] == 'a'
        assert mindb.get('task_args')[1] == 1
        assert mindb.get('task_args')[2].data == 67

        assert mindb.get('task_kwargs')['c'] == 6
        assert mindb.get('task_kwargs')['d'] == 'e'
        assert mindb.get('task_kwargs')['f'].data == 89

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        task_args = pickle.loads(b64decode(tr.task_args))
        assert task_args[0] == 'a'
        assert task_args[1] == 1
        assert task_args[2].data == 67

        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_kwargs['c'] == 6
        assert task_kwargs['d'] == 'e'
        assert task_kwargs['f'].data == 89

        tid3 = uuid()
        try:
            raise KeyError('foo')
        except KeyError as exception:
            self.b.mark_as_failure(tid3, exception)

        assert self.b.get_status(tid3) == states.FAILURE
        assert isinstance(self.b.get_result(tid3), KeyError)

    def test_backend__pickle_serialization__str_result__protocol_1(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, SomeClass(67)],
            kwargs={'c': 6, 'd': 'e', 'f': SomeClass(89)},
            task_protocol=1,
        )
        result = 'foo'

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result') == 'foo'
        assert mindb.get('task_name') == 'my_task'

        assert mindb.get('task_args')[0] == 'a'
        assert mindb.get('task_args')[1] == 1
        assert mindb.get('task_args')[2].data == 67

        assert mindb.get('task_kwargs')['c'] == 6
        assert mindb.get('task_kwargs')['d'] == 'e'
        assert mindb.get('task_kwargs')['f'].data == 89

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        task_args = pickle.loads(b64decode(tr.task_args))
        assert task_args[0] == 'a'
        assert task_args[1] == 1
        assert task_args[2].data == 67

        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_kwargs['c'] == 6
        assert task_kwargs['d'] == 'e'
        assert task_kwargs['f'].data == 89

    def test_backend__pickle_serialization__bytes_result__protocol_1(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, SomeClass(67)],
            kwargs={'c': 6, 'd': 'e', 'f': SomeClass(89)},
            task_protocol=1,
        )
        result = b'foo'

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result') == b'foo'
        assert mindb.get('task_name') == 'my_task'

        assert mindb.get('task_args')[0] == 'a'
        assert mindb.get('task_args')[1] == 1
        assert mindb.get('task_args')[2].data == 67

        assert mindb.get('task_kwargs')['c'] == 6
        assert mindb.get('task_kwargs')['d'] == 'e'
        assert mindb.get('task_kwargs')['f'].data == 89

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        task_args = pickle.loads(b64decode(tr.task_args))
        assert task_args[0] == 'a'
        assert task_args[1] == 1
        assert task_args[2].data == 67

        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_kwargs['c'] == 6
        assert task_kwargs['d'] == 'e'
        assert task_kwargs['f'].data == 89

    def test_backend__json_serialization__dict_result__protocol_1(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, True],
            kwargs={'c': 6, 'd': 'e', 'f': False},
            task_protocol=1,
        )
        result = {'foo': 'baz', 'bar': True}

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result').get('foo') == 'baz'
        assert mindb.get('result').get('bar') is True
        assert mindb.get('task_name') == 'my_task'
        assert mindb.get('task_args') == ['a', 1, True]
        assert mindb.get('task_kwargs') == {'c': 6, 'd': 'e', 'f': False}

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        assert json.loads(tr.task_args) == ['a', 1, True]
        assert json.loads(tr.task_kwargs) == {'c': 6, 'd': 'e', 'f': False}

        tid3 = uuid()
        try:
            raise KeyError('foo')
        except KeyError as exception:
            self.b.mark_as_failure(tid3, exception)

        assert self.b.get_status(tid3) == states.FAILURE
        assert isinstance(self.b.get_result(tid3), KeyError)

    def test_backend__json_serialization__str_result__protocol_1(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, True],
            kwargs={'c': 6, 'd': 'e', 'f': False},
            task_protocol=1,
        )
        result = 'foo'

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result') == 'foo'
        assert mindb.get('task_name') == 'my_task'
        assert mindb.get('task_args') == ['a', 1, True]
        assert mindb.get('task_kwargs') == {'c': 6, 'd': 'e', 'f': False}

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        assert json.loads(tr.task_args) == ['a', 1, True]
        assert json.loads(tr.task_kwargs) == {'c': 6, 'd': 'e', 'f': False}

    def test_backend__task_result_meta_injection(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=[],
            kwargs={},
            task_protocol=1,
        )
        result = None

        # inject request meta arbitrary data
        request.meta = {
            'key': 'value'
        }

        self.b.mark_as_done(tid2, result, request=request)
        mindb = self.b.get_task_meta(tid2)

        # check task meta
        assert mindb.get('result') is None
        assert mindb.get('task_name') == 'my_task'

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        assert json.loads(tr.meta) == {'key': 'value', 'children': []}

    def test_backend__task_result_date(self):
        tid2 = uuid()

        self.b.store_result(tid2, None, states.PENDING)

        tr = TaskResult.objects.get(task_id=tid2)
        assert tr.status == states.PENDING
        assert isinstance(tr.date_created, datetime.datetime)
        assert tr.date_started is None
        assert isinstance(tr.date_done, datetime.datetime)

        date_created = tr.date_created
        date_done = tr.date_done

        self.b.mark_as_started(tid2)

        tr = TaskResult.objects.get(task_id=tid2)
        assert tr.status == states.STARTED
        assert date_created == tr.date_created
        assert isinstance(tr.date_started, datetime.datetime)
        assert tr.date_done > date_done

        date_started = tr.date_started
        date_done = tr.date_done

        self.b.mark_as_done(tid2, 42)

        tr = TaskResult.objects.get(task_id=tid2)
        assert tr.status == states.SUCCESS
        assert tr.date_created == date_created
        assert tr.date_started == date_started
        assert tr.date_done > date_done

    def xxx_backend(self):
        tid = uuid()

        assert self.b.get_status(tid) == states.PENDING
        assert self.b.get_result(tid) is None

        self.b.mark_as_done(tid, 42)
        assert self.b.get_status(tid) == states.SUCCESS
        assert self.b.get_result(tid) == 42

        tid2 = uuid()
        try:
            raise KeyError('foo')
        except KeyError as exception:
            self.b.mark_as_failure(tid2, exception)

        assert self.b.get_status(tid2) == states.FAILURE
        assert isinstance(self.b.get_result(tid2), KeyError)

    def test_forget(self):
        tid = uuid()
        self.b.mark_as_done(tid, {'foo': 'bar'})
        x = self.app.AsyncResult(tid)
        assert x.result.get('foo') == 'bar'
        x.forget()
        if celery.VERSION[0:3] == (3, 1, 10):
            # bug in 3.1.10 means result did not clear cache after forget.
            x._cache = None
        assert x.result is None

    def test_secrets__pickle_serialization(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid = uuid()
        request = self._create_request(
            task_id=tid,
            name='my_task',
            args=['a', 1, 'password'],
            kwargs={'c': 3, 'd': 'e', 'password': 'password'},
            argsrepr='argsrepr',
            kwargsrepr='kwargsrepr',
        )
        result = {'foo': 'baz'}

        self.b.mark_as_done(tid, result, request=request)
        mindb = self.b.get_task_meta(tid)

        # check task meta
        assert mindb.get('result') == {'foo': 'baz'}
        assert mindb.get('task_args') == 'argsrepr'
        assert mindb.get('task_kwargs') == 'kwargsrepr'
        assert len(mindb.get('worker')) > 1

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid)
        task_args = pickle.loads(b64decode(tr.task_args))
        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_args == 'argsrepr'
        assert task_kwargs == 'kwargsrepr'

        # check async_result
        ar = AsyncResult(tid)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

    def test_secrets__json_serialization(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid = uuid()
        request = self._create_request(
            task_id=tid,
            name='my_task',
            args=['a', 1, True],
            kwargs={'c': 6, 'd': 'e', 'f': False},
            argsrepr='argsrepr',
            kwargsrepr='kwargsrepr',
        )
        result = {'foo': 'baz'}

        self.b.mark_as_done(tid, result, request=request)
        mindb = self.b.get_task_meta(tid)

        # check task meta
        assert mindb.get('result') == {'foo': 'baz'}
        assert mindb.get('task_args') == 'argsrepr'
        assert mindb.get('task_kwargs') == 'kwargsrepr'

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid)
        assert json.loads(tr.task_args) == 'argsrepr'
        assert json.loads(tr.task_kwargs) == 'kwargsrepr'

        # check async_result
        ar = AsyncResult(tid)
        assert ar.args == mindb.get('task_args')
        assert ar.kwargs == mindb.get('task_kwargs')

    def test_secrets__pickle_serialization__protocol_1(self):
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid = uuid()
        request = self._create_request(
            task_id=tid,
            name='my_task',
            args=['a', 1, SomeClass(67)],
            kwargs={'c': 6, 'd': 'e', 'f': SomeClass(89)},
            argsrepr='argsrepr',
            kwargsrepr='kwargsrepr',
            task_protocol=1,
        )
        result = {'foo': 'baz'}

        self.b.mark_as_done(tid, result, request=request)

        mindb = self.b.get_task_meta(tid)
        assert mindb.get('result') == {'foo': 'baz'}

        assert mindb.get('task_args')[0] == 'a'
        assert mindb.get('task_args')[1] == 1
        assert mindb.get('task_args')[2].data == 67

        assert mindb.get('task_kwargs')['c'] == 6
        assert mindb.get('task_kwargs')['d'] == 'e'
        assert mindb.get('task_kwargs')['f'].data == 89

        tr = TaskResult.objects.get(task_id=tid)
        task_args = pickle.loads(b64decode(tr.task_args))
        assert task_args[0] == 'a'
        assert task_args[1] == 1
        assert task_args[2].data == 67

        task_kwargs = pickle.loads(b64decode(tr.task_kwargs))
        assert task_kwargs['c'] == 6
        assert task_kwargs['d'] == 'e'
        assert task_kwargs['f'].data == 89

    def test_secrets__json_serialization__protocol_1(self):
        self.app.conf.result_serializer = 'json'
        self.app.conf.accept_content = {'pickle', 'json'}
        self.b = DatabaseBackend(app=self.app)

        tid = uuid()
        request = self._create_request(
            task_id=tid,
            name='my_task',
            args=['a', 1, True],
            kwargs={'c': 6, 'd': 'e', 'f': False},
            argsrepr='argsrepr',
            kwargsrepr='kwargsrepr',
            task_protocol=1,
        )
        result = {'foo': 'baz'}

        self.b.mark_as_done(tid, result, request=request)

        mindb = self.b.get_task_meta(tid)

        assert mindb.get('result') == {'foo': 'baz'}
        assert mindb.get('task_name') == 'my_task'
        assert mindb.get('task_args') == ['a', 1, True]
        assert mindb.get('task_kwargs') == {'c': 6, 'd': 'e', 'f': False}

        tr = TaskResult.objects.get(task_id=tid)
        assert json.loads(tr.task_args) == ['a', 1, True]
        assert json.loads(tr.task_kwargs) == {'c': 6, 'd': 'e', 'f': False}

    def test_apply_chord_header_result_arg(self):
        """Test if apply_chord can handle Celery <= 5.1 call signature"""
        gid = uuid()
        tid1 = uuid()
        tid2 = uuid()
        subtasks = [AsyncResult(tid1), AsyncResult(tid2)]
        group = GroupResult(id=gid, results=subtasks)
        # Celery < 5.1
        self.b.apply_chord(group, self.add.s())
        # Celery 5.1
        self.b.apply_chord((uuid(), subtasks), self.add.s())

    def test_on_chord_part_return(self):
        """Test if the ChordCounter is properly decremented and the callback is
        triggered after all chord parts have returned"""
        gid = uuid()
        tid1 = uuid()
        tid2 = uuid()
        subtasks = [AsyncResult(tid1), AsyncResult(tid2)]
        group = GroupResult(id=gid, results=subtasks)
        self.b.apply_chord(group, self.add.s())

        chord_counter = ChordCounter.objects.get(group_id=gid)
        assert chord_counter.count == 2

        request = mock.MagicMock()
        request.id = subtasks[0].id
        request.group = gid
        request.task = "my_task"
        request.args = ["a", 1, "password"]
        request.kwargs = {"c": 3, "d": "e", "password": "password"}
        request.argsrepr = "argsrepr"
        request.kwargsrepr = "kwargsrepr"
        request.hostname = "celery@ip-0-0-0-0"
        request.periodic_task_name = "my_periodic_task"
        request.ignore_result = False
        result = {"foo": "baz"}

        self.b.mark_as_done(tid1, result, request=request)

        chord_counter.refresh_from_db()
        assert chord_counter.count == 1

        self.b.mark_as_done(tid2, result, request=request)

        with pytest.raises(ChordCounter.DoesNotExist):
            ChordCounter.objects.get(group_id=gid)

        request.chord.delay.assert_called_once()

    def test_on_chord_part_return_counter_not_found(self):
        """Test if the chord does not raise an error if the ChordCounter is
        not found

        Basically this covers the case where a chord was created with a version
        <2.0.0 and the update was done before the chord was finished
        """
        request = mock.MagicMock()
        request.id = uuid()
        request.group = uuid()
        self.b.on_chord_part_return(request=request, state=None, result=None)

    def test_callback_failure(self):
        """Test if a failure in the chord callback is properly handled"""
        gid = uuid()
        tid1 = uuid()
        tid2 = uuid()
        cid = uuid()
        subtasks = [AsyncResult(tid1), AsyncResult(tid2)]
        group = GroupResult(id=gid, results=subtasks)
        self.b.apply_chord(group, self.add.s())

        chord_counter = ChordCounter.objects.get(group_id=gid)
        assert chord_counter.count == 2

        request = mock.MagicMock()
        request.id = subtasks[0].id
        request.group = gid
        request.task = "my_task"
        request.args = ["a", 1, "password"]
        request.kwargs = {"c": 3, "d": "e", "password": "password"}
        request.argsrepr = "argsrepr"
        request.kwargsrepr = "kwargsrepr"
        request.hostname = "celery@ip-0-0-0-0"
        request.periodic_task_name = "my_periodic_task"
        request.ignore_result = False
        request.chord.id = cid
        result = {"foo": "baz"}

        # Trigger an exception when the callback is triggered
        request.chord.delay.side_effect = ValueError()

        self.b.mark_as_done(tid1, result, request=request)

        chord_counter.refresh_from_db()
        assert chord_counter.count == 1

        self.b.mark_as_done(tid2, result, request=request)

        with pytest.raises(ChordCounter.DoesNotExist):
            ChordCounter.objects.get(group_id=gid)

        request.chord.delay.assert_called_once()

        assert TaskResult.objects.get(task_id=cid).status == states.FAILURE

    def test_on_chord_part_return_failure(self):
        """Test if a failure in one of the chord header tasks is properly
        handled and the callback was not triggered
        """
        gid = uuid()
        tid1 = uuid()
        tid2 = uuid()
        cid = uuid()
        subtasks = [AsyncResult(tid1), AsyncResult(tid2)]
        group = GroupResult(id=gid, results=subtasks)
        self.b.apply_chord(group, self.add.s())

        chord_counter = ChordCounter.objects.get(group_id=gid)
        assert chord_counter.count == 2

        request = mock.MagicMock()
        request.id = tid1
        request.group = gid
        request.task = "my_task"
        request.args = ["a", 1, "password"]
        request.kwargs = {"c": 3, "d": "e", "password": "password"}
        request.argsrepr = "argsrepr"
        request.kwargsrepr = "kwargsrepr"
        request.hostname = "celery@ip-0-0-0-0"
        request.periodic_task_name = "my_periodic_task"
        request.chord.id = cid
        result = {"foo": "baz"}

        self.b.mark_as_done(tid1, result, request=request)

        chord_counter.refresh_from_db()
        assert chord_counter.count == 1

        request.id = tid2
        self.b.mark_as_failure(tid2, ValueError(), request=request)

        with pytest.raises(ChordCounter.DoesNotExist):
            ChordCounter.objects.get(group_id=gid)

        request.chord.delay.assert_not_called()

    def test_groupresult_save_restore(self):
        """Test if we can save and restore a GroupResult"""
        group_id = uuid()
        results = [AsyncResult(id=uuid())]
        group = GroupResult(id=group_id, results=results)

        group.save(backend=self.b)

        restored_group = self.b.restore_group(group_id=group_id)

        assert restored_group == group

    def test_groupresult_save_restore_nested(self):
        """Test if we can save and restore a nested GroupResult"""
        group_id = uuid()
        async_result = AsyncResult(id=uuid())
        nested_results = [AsyncResult(id=uuid()), AsyncResult(id=uuid())]
        nested_group = GroupResult(id=uuid(), results=nested_results)
        group = GroupResult(id=group_id, results=[nested_group, async_result])

        group.save(backend=self.b)

        restored_group = self.b.restore_group(group_id=group_id)

        assert restored_group == group

    def test_backend_result_extended_is_false(self):
        self.app.conf.result_extended = False
        self.b = DatabaseBackend(app=self.app)
        tid2 = uuid()
        request = self._create_request(
            task_id=tid2,
            name='my_task',
            args=['a', 1, True],
            kwargs={'c': 6, 'd': 'e', 'f': False},
        )
        result = 'foo'

        self.b.mark_as_done(tid2, result, request=request)

        mindb = self.b.get_task_meta(tid2)

        # check meta data
        assert mindb.get('result') == 'foo'
        assert mindb.get('task_name') is None
        assert mindb.get('task_args') is None
        assert mindb.get('task_kwargs') is None

        # check task_result object
        tr = TaskResult.objects.get(task_id=tid2)
        assert tr.task_args is None
        assert tr.task_kwargs is None

    def test_custom_state(self):
        tid = uuid()

        assert self.b.get_status(tid) == states.PENDING

        self.b.store_result(tid, state="Progress", result={"progress": 10})
        assert self.b.get_status(tid) == "Progress"
        assert self.b.get_result(tid) == {"progress": 10}

        self.b.mark_as_done(tid, 42)
        assert self.b.get_status(tid) == states.SUCCESS
        assert self.b.get_result(tid) == 42


class DjangoCeleryResultRouter:
    route_app_labels = {"django_celery_results"}

    def db_for_read(self, model, **hints):
        """Route read access to the read-only database"""
        if model._meta.app_label in self.route_app_labels:
            return "read-only"
        return None

    def db_for_write(self, model, **hints):
        """Route write access to the default database"""
        if model._meta.app_label in self.route_app_labels:
            return "default"
        return None


class ChordPartReturnTestCase(TransactionTestCase):
    databases = {"default", "read-only"}

    def setUp(self):
        super().setUp()
        self.app.conf.result_serializer = 'json'
        self.app.conf.result_backend = (
            'django_celery_results.backends:DatabaseBackend')
        self.app.conf.result_extended = True
        self.b = DatabaseBackend(app=self.app)

    def test_on_chord_part_return_multiple_databases(self):
        """
        Test if the ChordCounter is properly decremented and the callback is
        triggered after all chord parts have returned with multiple databases
        """
        with self.settings(DATABASE_ROUTERS=[DjangoCeleryResultRouter()]):
            gid = uuid()
            tid1 = uuid()
            tid2 = uuid()
            subtasks = [AsyncResult(tid1), AsyncResult(tid2)]
            group = GroupResult(id=gid, results=subtasks)

            assert ChordCounter.objects.count() == 0
            assert ChordCounter.objects.using("read-only").count() == 0
            assert ChordCounter.objects.using("default").count() == 0

            self.b.apply_chord(group, self.add.s())

            # Check if the ChordCounter was created in the correct database
            assert ChordCounter.objects.count() == 1
            assert ChordCounter.objects.using("read-only").count() == 1
            assert ChordCounter.objects.using("default").count() == 1

            chord_counter = ChordCounter.objects.get(group_id=gid)
            assert chord_counter.count == 2

            request = mock.MagicMock()
            request.id = subtasks[0].id
            request.group = gid
            request.task = "my_task"
            request.args = ["a", 1, "password"]
            request.kwargs = {"c": 3, "d": "e", "password": "password"}
            request.argsrepr = "argsrepr"
            request.kwargsrepr = "kwargsrepr"
            request.hostname = "celery@ip-0-0-0-0"
            request.periodic_task_name = "my_periodic_task"
            request.ignore_result = False
            result = {"foo": "baz"}

            self.b.mark_as_done(tid1, result, request=request)

            chord_counter.refresh_from_db()
            assert chord_counter.count == 1

            self.b.mark_as_done(tid2, result, request=request)

            with pytest.raises(ChordCounter.DoesNotExist):
                ChordCounter.objects.get(group_id=gid)

            request.chord.delay.assert_called_once()