# 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.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 BaseTaskData, TaskTypes
from debusine.test.django import TestCase, TransactionTestCase


class TestTaskNotManagingTransactions(BaseServerTask[BaseTaskData]):
    """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(BaseServerTask[BaseTaskData]):
    """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.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.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."""
        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()
        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)
