# Copyright 2024 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Unit tests for Celery integration."""

import logging
from unittest import mock

from django.db import connections, transaction

from debusine.db.context import context
from debusine.db.models import WorkRequest, Worker
from debusine.server.celery import (
    CeleryWorkerRequiresServerTask,
    run_server_task,
)
from debusine.server.tasks import BaseServerTask
from debusine.tasks import TaskConfigError
from debusine.tasks.models import BaseDynamicTaskData, BaseTaskData, TaskTypes
from debusine.tasks.tests.helper_mixin import TestTaskMixin
from debusine.test.django import TestCase, TransactionTestCase


class TestTaskNotManagingTransactions(
    TestTaskMixin, BaseServerTask[BaseTaskData, BaseDynamicTaskData]
):
    """A sample server task that does not manage its own transactions."""

    TASK_VERSION = 1

    def _execute(self) -> bool:
        return connections["default"].in_atomic_block


class TestTaskManagingTransactions(
    TestTaskMixin, BaseServerTask[BaseTaskData, BaseDynamicTaskData]
):
    """A sample server task that manages its own transactions."""

    TASK_VERSION = 1
    TASK_MANAGES_TRANSACTIONS = True

    def _execute(self) -> bool:
        return not connections["default"].in_atomic_block


class RunServerTaskTests(TestCase):
    """Unit tests for :py:func:`run_server_task`."""

    def setUp(self):
        """Create common objects."""
        self.worker = Worker.objects.get_or_create_celery()

    def create_assigned_work_request(self, **kwargs):
        """Create a work request and assign it to self.worker."""
        work_request = self.playground.create_work_request(**kwargs)
        work_request.assign_worker(self.worker)
        return work_request

    def test_no_work_request(self):
        """The Celery task fails if the work request does not exist."""
        work_request_id = self.playground.create_work_request().id + 1
        with self.assertLogsContains(
            f"Work request {work_request_id} does not exist",
            logger="debusine.server.celery",
            level=logging.ERROR,
        ):
            result = run_server_task.apply(args=(work_request_id,))
        self.assertTrue(result.failed())
        self.assertIsInstance(result.result, WorkRequest.DoesNotExist)

    def test_bad_task_name(self):
        """The Celery task fails if the work request has a bad task name."""
        work_request = self.create_assigned_work_request(
            task_type=TaskTypes.SERVER, task_name="nonexistent"
        )
        with (
            self.assertLogsContains(
                "Task: nonexistent does not exist",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertTrue(result.failed())
        self.assertIsInstance(result.result, ValueError)
        work_request.refresh_from_db()
        self.assertIsNotNone(work_request.started_at)
        self.assertIsNotNone(work_request.completed_at)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.ERROR)

    def test_bad_task_data(self):
        """The Celery task fails if the work request has bad task data."""
        work_request = self.create_assigned_work_request(
            task_type=TaskTypes.SERVER,
            task_name="servernoop",
            task_data={"nonexistent": True},
        )
        with (
            self.assertLogsContains(
                "Task: servernoop failed to configure",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertTrue(result.failed())
        self.assertIsInstance(result.result, TaskConfigError)
        work_request.refresh_from_db()
        self.assertIsNotNone(work_request.started_at)
        self.assertIsNotNone(work_request.completed_at)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.ERROR)

    def test_external_task(self):
        """The Celery task fails if asked to run an external task."""
        work_request = self.create_assigned_work_request(
            task_name="noop", task_data={"result": True}
        )
        with (
            self.assertLogsContains(
                "Task: noop cannot run on a Celery worker",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertTrue(result.failed())
        self.assertIsInstance(result.result, CeleryWorkerRequiresServerTask)
        work_request.refresh_from_db()
        self.assertIsNotNone(work_request.started_at)
        self.assertIsNotNone(work_request.completed_at)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.ERROR)

    def test_task_raises_exception(self):
        """The Celery task fails if the task raises an exception."""
        work_request = self.create_assigned_work_request(
            task_type=TaskTypes.SERVER,
            task_name="servernoop",
            task_data={"exception": True},
        )
        with (
            self.assertLogsContains(
                "Task: servernoop failed to execute",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertTrue(result.failed())
        self.assertIsInstance(result.result, RuntimeError)
        self.assertEqual(result.result.args, ("Client requested an exception",))
        work_request.refresh_from_db()
        self.assertIsNotNone(work_request.started_at)
        self.assertIsNotNone(work_request.completed_at)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.ERROR)

    def test_task_aborted(self):
        """If the task is aborted, the Celery task logs but does not fail."""
        work_request = self.create_assigned_work_request(
            task_type=TaskTypes.SERVER, task_name="servernoop"
        )
        with (
            self.assertLogsContains(
                "Task: servernoop has been aborted",
                logger="debusine.server.celery",
                level=logging.INFO,
            ),
            mock.patch(
                "debusine.server.tasks.noop.ServerNoop.execute",
                lambda task: task.abort(),
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertFalse(result.failed())
        self.assertFalse(result.result)
        work_request.refresh_from_db()
        self.assertIsNotNone(work_request.started_at)
        self.assertIsNone(work_request.completed_at)
        self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(work_request.result, WorkRequest.Results.NONE)

    def test_task_returns_false(self):
        """If the task returns False, the Celery task records that."""
        work_request = self.create_assigned_work_request(
            task_type=TaskTypes.SERVER,
            task_name="servernoop",
            task_data={"result": False},
        )
        with transaction.atomic():
            result = run_server_task.apply(args=(work_request.id,))
        self.assertFalse(result.failed())
        self.assertFalse(result.result)
        work_request.refresh_from_db()
        self.assertIsNotNone(work_request.started_at)
        self.assertIsNotNone(work_request.completed_at)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.FAILURE)

    def test_task_returns_true(self):
        """If the task returns True, the Celery task records that."""
        work_request = self.create_assigned_work_request(
            task_type=TaskTypes.SERVER,
            task_name="servernoop",
            task_data={"result": True},
        )
        with transaction.atomic():
            result = run_server_task.apply(args=(work_request.id,))
        self.assertFalse(result.failed())
        self.assertTrue(result.result)
        work_request.refresh_from_db()
        self.assertIsNotNone(work_request.started_at)
        self.assertIsNotNone(work_request.completed_at)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.SUCCESS)


class RunServerTaskTransactionTests(TransactionTestCase):
    """Unit tests for :py:func:`run_server_task`."""

    def test_task_not_managing_transactions(self):
        """An atomic block is used for tasks that do not manage transactions."""
        with context.disable_permission_checks():
            worker = Worker.objects.get_or_create_celery()
            work_request = self.playground.create_work_request(
                task_type=TaskTypes.SERVER,
                task_name="testtasknotmanagingtransactions",
                task_data={},
            )
            work_request.assign_worker(worker)
        result = run_server_task.apply(args=(work_request.id,))
        self.assertFalse(result.failed())
        self.assertTrue(result.result)

    def test_task_managing_transactions(self):
        """An atomic block is not used for tasks that manage transactions."""
        worker = Worker.objects.get_or_create_celery()
        with context.disable_permission_checks():
            work_request = self.playground.create_work_request(
                task_type=TaskTypes.SERVER,
                task_name="testtaskmanagingtransactions",
                task_data={},
            )
        work_request.assign_worker(worker)
        result = run_server_task.apply(args=(work_request.id,))
        self.assertFalse(result.failed())
        self.assertTrue(result.result)
