# Copyright © 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 typing import Any, override
from unittest import mock

from django.db import connections, transaction

from debusine.db.models import WorkRequest, Worker
from debusine.db.models.work_requests import workflow_ancestors
from debusine.db.tests.utils import RunInThreadAndCloseDBConnections
from debusine.server.celery import (
    WorkRequestNotPending,
    run_server_task,
    update_workflows,
)
from debusine.server.tasks.models import ServerNoopData
from debusine.server.tasks.tests.helpers import SampleBaseServerTask
from debusine.tasks.models import (
    BaseDynamicTaskData,
    BaseTaskData,
    OutputData,
    OutputDataError,
)
from debusine.test.django import TestCase, TransactionTestCase


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

    TASK_VERSION = 1

    @override
    def _execute(self) -> WorkRequest.Results:
        return (
            WorkRequest.Results.SUCCESS
            if connections["default"].in_atomic_block
            else WorkRequest.Results.FAILURE
        )


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

    TASK_VERSION = 1
    TASK_MANAGES_TRANSACTIONS = True

    @override
    def _execute(self) -> WorkRequest.Results:
        return (
            WorkRequest.Results.SUCCESS
            if not connections["default"].in_atomic_block
            else WorkRequest.Results.FAILURE
        )


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

    def setUp(self) -> None:
        """Create common objects."""
        super().setUp()
        self.worker = Worker.objects.get_or_create_celery()

    def create_assigned_work_request(
        self, task_name: str = "servernoop", **kwargs: Any
    ) -> WorkRequest:
        """Create a work request and assign it to self.worker."""
        work_request = self.playground.create_server_task(
            task_name=task_name, mark_pending=True, **kwargs
        )
        work_request.assign_worker(self.worker)
        return work_request

    def test_no_work_request(self) -> None:
        """The Celery task fails if the work request does not exist."""
        work_request_id = self.playground.create_worker_task().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,))
        with self.assertRaises(WorkRequest.DoesNotExist):
            result.get()

    def test_not_pending(self) -> None:
        """The Celery task fails if the work request is not pending."""
        work_request = self.create_assigned_work_request(
            result=WorkRequest.Results.ERROR
        )
        with self.assertLogsContains(
            f"Work request {work_request.id} is in status completed, not "
            f"pending",
            logger="debusine.server.celery",
            level=logging.ERROR,
        ):
            result = run_server_task.apply(args=(work_request.id,))
        with self.assertRaises(WorkRequestNotPending):
            result.get()

    def test_bad_task_name(self) -> None:
        """The Celery task logs but does not fail if given a bad task name."""
        work_request = self.create_assigned_work_request()
        work_request.task_name = "nonexistent"
        work_request.save()
        expected_message = "'nonexistent' is not a registered Server task_name"
        with (
            self.assertLogsContains(
                f"Error running work request Server/nonexistent "
                f"({work_request.id}): {expected_message}",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.ERROR)
        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)
        self.assertEqual(
            work_request.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="setup-failed"
                    )
                ]
            ),
        )

    def test_bad_task_data(self) -> None:
        """The Celery task logs but does not fail if given bad task data."""
        work_request = self.create_assigned_work_request()
        work_request.task_data = {"nonexistent": True}
        work_request.configured_task_data = {"nonexistent": True}
        work_request.save()
        expected_message = (
            "Failed to configure: 1 validation error for ServerNoopData"
        )
        with (
            self.assertLogsContains(
                f"Error running work request Server/servernoop "
                f"({work_request.id}): {expected_message}",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.ERROR)
        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)
        assert work_request.output_data is not None
        self.assertIsInstance(work_request.output_data, OutputData)
        assert work_request.output_data.errors is not None
        self.assertEqual(len(work_request.output_data.errors), 1)
        self.assertTrue(
            work_request.output_data.errors[0].message.startswith(
                expected_message
            )
        )
        self.assertEqual(
            work_request.output_data.errors[0].code, "configure-failed"
        )

    def test_external_task(self) -> None:
        """The Celery task logs but does not fail on an external task."""
        work_request = self.playground.create_worker_task(
            task_name="noop",
            task_data={"result": WorkRequest.Results.SUCCESS},
            mark_pending=True,
            worker=self.worker,
        )
        expected_message = "Cannot run on a Celery worker"
        with (
            self.assertLogsContains(
                f"Error running work request Worker/noop ({work_request.id}): "
                f"{expected_message}",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.ERROR)
        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)
        self.assertEqual(
            work_request.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="wrong-task-type"
                    )
                ]
            ),
        )

    def test_task_raises_exception(self) -> None:
        """The Celery task logs but does not fail if the task raises."""
        work_request = self.create_assigned_work_request(
            task_data=ServerNoopData(exception=True),
        )
        expected_message = "Execution failed: Client requested an exception"
        with (
            self.assertLogsContains(
                f"Error running work request Server/servernoop "
                f"({work_request.id}): {expected_message}",
                logger="debusine.server.celery",
                level=logging.ERROR,
            ),
            transaction.atomic(),
        ):
            result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.ERROR)
        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)
        self.assertEqual(
            work_request.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="execute-failed"
                    )
                ]
            ),
        )

    def test_task_aborted(self) -> None:
        """If the task is aborted, the Celery task logs but does not fail."""
        work_request = self.create_assigned_work_request()
        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.assertEqual(result.get(), WorkRequest.Results.FAILURE)
        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)
        self.assertIsNone(work_request.output_data)

    def test_task_returns_failure(self) -> None:
        """If the task returns FAILURE, the Celery task records that."""
        work_request = self.create_assigned_work_request(
            task_data=ServerNoopData(result=WorkRequest.Results.FAILURE),
        )
        with transaction.atomic():
            result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.FAILURE)
        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)
        self.assertIsNone(work_request.output_data)

    def test_task_returns_true(self) -> None:
        """If the task returns True, the Celery task records that."""
        work_request = self.create_assigned_work_request(
            task_data=ServerNoopData(result=WorkRequest.Results.SUCCESS),
        )
        with transaction.atomic():
            result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.SUCCESS)
        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)
        self.assertIsNone(work_request.output_data)


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

    def test_task_not_managing_transactions(self) -> None:
        """An atomic block is used for tasks that do not manage transactions."""
        worker = Worker.objects.get_or_create_celery()
        work_request = self.playground.create_server_task(
            task_name="sampletasknotmanagingtransactions",
            task_data={},
            mark_pending=True,
        )
        work_request.assign_worker(worker)
        result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.SUCCESS)

    def test_task_managing_transactions(self) -> None:
        """An atomic block is not used for tasks that manage transactions."""
        worker = Worker.objects.get_or_create_celery()
        work_request = self.playground.create_server_task(
            task_name="sampletaskmanagingtransactions",
            task_data={},
            mark_pending=True,
        )
        work_request.assign_worker(worker)
        result = run_server_task.apply(args=(work_request.id,))
        self.assertEqual(result.get(), WorkRequest.Results.SUCCESS)

    def test_locks_workflow_before_completion(self) -> None:
        workflow = self.playground.create_workflow()
        server_noop = workflow.create_child_server(
            task_name="servernoop",
            task_data=ServerNoopData(result=WorkRequest.Results.SUCCESS),
        )
        noop = workflow.create_child_worker(task_name="noop")
        self.playground.advance_work_request(server_noop, mark_pending=True)
        self.playground.assign_worker(server_noop)
        unrelated = self.playground.create_worker_task()

        with mock.patch.object(
            WorkRequest,
            "mark_completed",
            side_effect=(
                lambda *args, **kwargs: self.assert_work_requests_locked(
                    locked=[workflow, server_noop, noop], unlocked=[unrelated]
                )
            ),
        ):
            result = run_server_task.apply(args=(server_noop.id,))

        self.assertEqual(result.get(), WorkRequest.Results.SUCCESS)

    def test_locks_workflow_before_completion_error(self) -> None:
        workflow = self.playground.create_workflow()
        server_noop = workflow.create_child_server(
            task_name="servernoop", task_data=ServerNoopData(exception=True)
        )
        noop = workflow.create_child_worker(task_name="noop")
        self.playground.advance_work_request(server_noop, mark_pending=True)
        self.playground.assign_worker(server_noop)
        unrelated = self.playground.create_worker_task()

        with mock.patch.object(
            WorkRequest,
            "mark_completed",
            side_effect=(
                lambda *args, **kwargs: self.assert_work_requests_locked(
                    locked=[workflow, server_noop, noop], unlocked=[unrelated]
                )
            ),
        ):
            result = run_server_task.apply(args=(server_noop.id,))

        self.assertEqual(result.get(), WorkRequest.Results.ERROR)


class UpdateWorkflowTransactionTests(TransactionTestCase):
    """Unit tests for :py:func:`update_workflows`."""

    def test_updates_workflows(self) -> None:
        """The Celery task updates workflow properties."""
        with (
            transaction.atomic(),
            mock.patch("django.db.transaction.on_commit"),
        ):
            workflow = self.playground.create_workflow()
            work_request = workflow.create_child_worker("noop")
            self.assertFalse(work_request.workflows_need_update)
            work_request.mark_aborted()
            self.assertTrue(work_request.workflows_need_update)
            unrelated = self.playground.create_worker_task()

        with transaction.atomic():
            result = update_workflows.apply()
            self.assertIsNone(result.get())

            workflow.refresh_from_db()
            work_request.refresh_from_db()
            self.assertEqual(
                workflow.workflow_last_activity_at, work_request.completed_at
            )
            self.assertEqual(
                workflow.workflow_runtime_status,
                WorkRequest.RuntimeStatuses.ABORTED,
            )
            self.assertFalse(work_request.workflows_need_update)
            self.assert_work_requests_locked(
                locked=[workflow, work_request], unlocked=[unrelated]
            )

    def test_repeatable_read(self) -> None:
        """The task runs in the correct transaction isolation level."""
        with (
            transaction.atomic(),
            mock.patch("django.db.transaction.on_commit"),
        ):
            workflows = [self.playground.create_workflow() for _ in range(2)]
            work_requests = [
                workflow.create_child_worker("noop") for workflow in workflows
            ]
            self.assertFalse(work_requests[0].workflows_need_update)
            work_requests[0].mark_aborted()
            self.assertTrue(work_requests[0].workflows_need_update)
            self.assertFalse(work_requests[1].workflows_need_update)

        def create_workflow_needs_update() -> None:
            with transaction.atomic():
                with mock.patch("django.db.transaction.on_commit"):
                    work_requests[1].mark_aborted()
                self.assertTrue(work_requests[1].workflows_need_update)

        def create_workflow_needs_update_in_thread(
            *args: Any, **kwargs: Any
        ) -> Any:
            thread = RunInThreadAndCloseDBConnections(
                create_workflow_needs_update
            )
            thread.run_and_wait()
            return mock.DEFAULT

        with mock.patch(
            "debusine.server.celery.workflow_ancestors",
            side_effect=create_workflow_needs_update_in_thread,
            wraps=workflow_ancestors,
        ):
            result = update_workflows.apply()
            self.assertIsNone(result.get())

        # The first workflow is updated.
        workflows[0].refresh_from_db()
        work_requests[0].refresh_from_db()
        self.assertEqual(
            workflows[0].workflow_last_activity_at,
            work_requests[0].completed_at,
        )
        self.assertFalse(work_requests[0].workflows_need_update)

        # The second workflow is untouched, since it didn't need an update
        # until after our task's transaction started.
        workflows[1].refresh_from_db()
        work_requests[1].refresh_from_db()
        self.assertIsNone(workflows[1].workflow_last_activity_at)
        self.assertTrue(work_requests[1].workflows_need_update)
