# 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 Debusine Celery configuration."""

import json
from multiprocessing import Lock
from pathlib import Path
from typing import Any
from unittest import mock

from celery import shared_task, signals, states
from celery.contrib.testing.worker import start_worker
from celery.utils.nodenames import gethostname, nodename
from django.conf import settings
from django.db import transaction
from django.test import override_settings
from django_celery_results.models import TaskResult

import debusine.worker.tags as wtags
from debusine.db.models import Token, Worker
from debusine.project.celery import (
    CeleryToken,
    app,
    apply_async_on_commit,
    get_celery_token,
    get_celery_worker_metadata,
    get_celery_worker_provided_tags,
    get_scheduler_worker_metadata,
    make_app,
)
from debusine.tasks.models import WorkerType
from debusine.test.django import TestCase, TransactionTestCase

_lock = Lock()


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task  # type: ignore[untyped-decorator]
def _locked_noop() -> None:
    """Sample Celery task to do nothing, once a lock is available."""
    with _lock:
        pass


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task  # type: ignore[untyped-decorator]
def _locked_reverse_string(argument: str, **kwargs: Any) -> str:
    """Sample Celery task to reverse a string, once a lock is available."""
    with _lock:
        return "".join(reversed(argument))


@override_settings(CELERY_BROKER_URL="memory://")
class TestApplyAsyncOnCommit(TransactionTestCase):
    """Tests for ``apply_async_on_commit`` and ``store_pending_state``."""

    def setUp(self) -> None:
        super().setUp()
        self.app = make_app()

    def test_result_states(self) -> None:
        """The task's state can be tracked in the database."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(self.app, perform_ping_check=False) as celery_worker,
        ):
            expected_task_name = f"{__name__}._locked_reverse_string"
            expected_task_args = ["hello"]
            expected_task_kwargs = {"foo": "bar", "value": 1}

            def assert_task_ok(task_id: str, expected_status: str) -> None:
                task_result = TaskResult.objects.get(task_id=task_id)
                self.assertEqual(task_result.status, expected_status)
                self.assertEqual(task_result.task_name, expected_task_name)
                self.assertEqual(
                    json.loads(task_result.task_args), expected_task_args
                )
                self.assertEqual(
                    json.loads(task_result.task_kwargs), expected_task_kwargs
                )

            # Schedule a task, but with a lock held so that we get a chance
            # to look at its result in the database before it finishes.
            with _lock:
                with transaction.atomic():
                    task_id = apply_async_on_commit(
                        _locked_reverse_string,
                        args=("hello",),
                        kwargs={"foo": "bar", "value": 1},
                    )
                    assert task_id is not None

                    # The result before committing the transaction is correct.
                    assert_task_ok(task_id, states.PENDING)

                # The result after committing the transaction is correct.
                assert_task_ok(task_id, states.PENDING)

            # Release the lock to allow the task to complete.  Its result in
            # the database is updated with the new status.
            self.assertEqual(
                _locked_reverse_string.AsyncResult(task_id).get(), "olleh"
            )

            assert_task_ok(task_id, states.SUCCESS)

            signals.worker_shutdown.send(sender=celery_worker)

    def test_result_states_no_arguments(self) -> None:
        """A task with no arguments is tracked correctly in the database."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(self.app, perform_ping_check=False) as celery_worker,
        ):
            expected_task_name = f"{__name__}._locked_noop"

            def assert_task_ok(task_id: str, expected_status: str) -> None:
                task_result = TaskResult.objects.get(task_id=task_id)
                self.assertEqual(task_result.status, expected_status)
                self.assertEqual(task_result.task_name, expected_task_name)
                self.assertEqual(json.loads(task_result.task_args), [])
                self.assertEqual(json.loads(task_result.task_kwargs), {})

            # Schedule a task, but with a lock held so that we get a chance
            # to look at its result in the database before it finishes.
            with _lock:
                with transaction.atomic():
                    task_id = apply_async_on_commit(_locked_noop)
                    assert task_id is not None

                    # The result before committing the transaction is correct.
                    assert_task_ok(task_id, states.PENDING)

                # The result after committing the transaction is correct.
                assert_task_ok(task_id, states.PENDING)

            # Release the lock to allow the task to complete.  Its result in
            # the database is updated with the new status.
            self.assertIsNone(_locked_noop.AsyncResult(task_id).get())

            assert_task_ok(task_id, states.SUCCESS)

            signals.worker_shutdown.send(sender=celery_worker)

    def test_result_states_once_per_savepoint(self) -> None:
        """``once_per_savepoint=True`` allows batching up tasks."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(self.app, perform_ping_check=False) as celery_worker,
        ):
            expected_task_name = f"{__name__}._locked_noop"

            def assert_task_ok(task_id: str, expected_status: str) -> None:
                task_result = TaskResult.objects.get(task_id=task_id)
                self.assertEqual(task_result.status, expected_status)
                self.assertEqual(task_result.task_name, expected_task_name)
                self.assertEqual(json.loads(task_result.task_args), [])
                self.assertEqual(json.loads(task_result.task_kwargs), {})

            with _lock:
                with transaction.atomic():
                    with transaction.atomic():
                        task_id = apply_async_on_commit(
                            _locked_noop, once_per_savepoint=True
                        )
                        assert task_id is not None
                        assert_task_ok(task_id, states.PENDING)

                        # If the task has already been scheduled in this
                        # savepoint, no new task is scheduled.
                        self.assertIsNone(
                            apply_async_on_commit(
                                _locked_noop, once_per_savepoint=True
                            )
                        )

                    # If the task has not been scheduled in this savepoint,
                    # a new task is scheduled.
                    task_id_2 = apply_async_on_commit(
                        _locked_noop, once_per_savepoint=True
                    )
                    assert task_id_2 is not None
                    self.assertNotEqual(task_id_2, task_id)
                    assert_task_ok(task_id_2, states.PENDING)

                assert_task_ok(task_id, states.PENDING)
                assert_task_ok(task_id_2, states.PENDING)

            self.assertIsNone(_locked_noop.AsyncResult(task_id).get())
            self.assertIsNone(_locked_noop.AsyncResult(task_id_2).get())

            assert_task_ok(task_id, states.SUCCESS)
            assert_task_ok(task_id_2, states.SUCCESS)

            signals.worker_shutdown.send(sender=celery_worker)


class TestCeleryToken(TestCase):
    """Tests for :py:class:`CeleryToken`."""

    def test_existing_token(self) -> None:
        tokens_dir = self.create_temporary_directory()
        key = "existing-key"
        Path(tokens_dir, "test").write_text(key)

        with override_settings(
            DEBUSINE_FQDN="debusine.example.org",
            DEBUSINE_TOKENS_DIRECTORY=str(tokens_dir),
        ):
            celery_token = CeleryToken("test")
        self.addCleanup(celery_token.token.delete)

        self.assertEqual(celery_token.token.hash, Token._generate_hash(key))
        self.assertEqual(
            celery_token.client.base_api_url, "https://debusine.example.org/api"
        )
        self.assertEqual(celery_token.client.token, key)
        self.assertIsNone(celery_token.client.scope)

    def test_new_token(self) -> None:
        tokens_dir = self.create_temporary_directory()

        with override_settings(
            DEBUSINE_FQDN="debusine.example.org",
            DEBUSINE_TOKENS_DIRECTORY=str(tokens_dir),
        ):
            celery_token = CeleryToken("test")
        self.addCleanup(celery_token.token.delete)

        key = Path(tokens_dir, "test").read_text()
        self.assertEqual(len(key), 64)
        self.assertEqual(celery_token.token.hash, Token._generate_hash(key))
        self.assertEqual(
            celery_token.client.base_api_url, "https://debusine.example.org/api"
        )
        self.assertEqual(celery_token.client.token, key)
        self.assertIsNone(celery_token.client.scope)


class TestWorkerMetadata(TestCase):
    """Test for worker metadata generation functions."""

    def test_celery_tags(self) -> None:
        system_provided_tags = self.enterContext(
            self.spy("debusine.worker.system_information.system_provided_tags")
        )
        get_provided_celery_worker_tags = self.enterContext(
            mock.patch(
                "debusine.server.tasks.get_provided_celery_worker_tags",
                return_value={"servertasks"},
            )
        )
        executor_provided_tags = self.enterContext(
            mock.patch(
                "debusine.tasks.executors.executor_provided_tags",
                return_value={"executor"},
            )
        )
        tags = get_celery_worker_provided_tags()
        self.assertIn(wtags.WORKER_TYPE_CELERY, tags)
        self.assertNotIn("executor", tags)
        self.assertIn("servertasks", tags)

        system_provided_tags.assert_called_once_with(WorkerType.CELERY)
        executor_provided_tags.assert_not_called()
        get_provided_celery_worker_tags.assert_called_once_with()

    def test_celery(self) -> None:
        system_metadata = self.enterContext(
            self.spy("debusine.worker.system_information.system_metadata")
        )
        analyze_celery_worker_tasks = self.enterContext(
            self.spy("debusine.server.tasks.analyze_celery_worker_tasks")
        )
        analyze_worker_all_executors = self.enterContext(
            self.spy("debusine.tasks.executors.analyze_worker_all_executors")
        )
        get_celery_worker_provided_tags = self.enterContext(
            self.spy("debusine.project.celery.get_celery_worker_provided_tags")
        )

        md = get_celery_worker_metadata()
        self.assertEqual(md["system:worker_type"], WorkerType.CELERY)
        self.assertIn("provided_tags", md)

        system_metadata.assert_called_once_with(WorkerType.CELERY)
        analyze_celery_worker_tasks.assert_called_once_with()
        analyze_worker_all_executors.assert_not_called()
        get_celery_worker_provided_tags.assert_called_once_with()

    def test_scheduler(self) -> None:
        system_metadata = self.enterContext(
            self.spy("debusine.worker.system_information.system_metadata")
        )
        md = get_scheduler_worker_metadata()
        self.assertEqual(md["system:worker_type"], WorkerType.SCHEDULER)
        self.assertNotIn("provided_tags", md)

        system_metadata.assert_called_once_with(WorkerType.SCHEDULER)


@override_settings(CELERY_BROKER_URL="memory://")
class TestConnectWorker(TransactionTestCase):
    """Tests for :py:func:`connect_worker`."""

    def setUp(self) -> None:
        super().setUp()
        self.app = make_app()
        Worker.objects.filter(
            worker_type__in={WorkerType.CELERY, WorkerType.SCHEDULER}
        ).delete()

    def test_server_task_worker(self) -> None:
        """The main worker creates a Worker instance that runs server tasks."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            # Set a concurrency that's unlikely to be the actual number of
            # CPUs we have.
            start_worker(
                self.app,
                concurrency=5,
                perform_ping_check=False,
                hostname=nodename("celery", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.CELERY)
            self.assertEqual(worker.concurrency, 5)
            self.assertEqual(
                worker.dynamic_metadata["system:worker_type"], WorkerType.CELERY
            )
            self.assertIn("provided_tags", worker.dynamic_metadata)
            self.assertIsNotNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.SCHEDULER).exists()
            )

            token = Token.objects.get_token_or_none(
                Path(settings.DEBUSINE_TOKENS_DIRECTORY, "celery").read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

            worker.refresh_from_db()
            self.assertIsNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.SCHEDULER).exists()
            )

    def test_server_task_worker_token_already_exists(self) -> None:
        with transaction.atomic():
            existing_worker = Worker.objects.get_or_create_celery()
            existing_worker.token = CeleryToken("celery").token
            existing_worker.save()

        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            # Set a concurrency that's unlikely to be the actual number of
            # CPUs we have.
            start_worker(
                self.app,
                concurrency=5,
                perform_ping_check=False,
                hostname=nodename("celery", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.CELERY)
            self.assertEqual(worker, existing_worker)

            token = Token.objects.get_token_or_none(
                Path(settings.DEBUSINE_TOKENS_DIRECTORY, "celery").read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            self.assertEqual(token, existing_worker.token)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

    def test_scheduler_worker(self) -> None:
        """The scheduler worker creates a Worker instance."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(
                app,
                concurrency=1,
                perform_ping_check=False,
                hostname=nodename("scheduler", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.SCHEDULER)
            self.assertEqual(worker.concurrency, 1)
            self.assertEqual(
                worker.dynamic_metadata["system:worker_type"],
                WorkerType.SCHEDULER,
            )
            self.assertIsNotNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.CELERY).exists()
            )

            token = Token.objects.get_token_or_none(
                Path(
                    settings.DEBUSINE_TOKENS_DIRECTORY, "scheduler"
                ).read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

            worker.refresh_from_db()
            self.assertIsNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.CELERY).exists()
            )

    def test_scheduler_worker_token_already_exists(self) -> None:
        with transaction.atomic():
            existing_worker = Worker.objects.get_or_create_scheduler()
            existing_worker.token = CeleryToken("scheduler").token
            existing_worker.save()

        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(
                app,
                concurrency=1,
                perform_ping_check=False,
                hostname=nodename("scheduler", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.SCHEDULER)
            self.assertEqual(worker, existing_worker)

            token = Token.objects.get_token_or_none(
                Path(
                    settings.DEBUSINE_TOKENS_DIRECTORY, "scheduler"
                ).read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            self.assertEqual(token, existing_worker.token)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

    def test_provisioner_worker(self) -> None:
        """The provisioner worker does not create a Worker instance."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(
                app,
                concurrency=1,
                perform_ping_check=False,
                hostname=nodename("provisioner", gethostname()),
            ) as celery_worker,
        ):
            self.assertFalse(
                Worker.objects.filter(
                    worker_type__in={WorkerType.CELERY, WorkerType.SCHEDULER}
                ).exists()
            )
            self.assertIsNone(get_celery_token())

            signals.worker_shutdown.send(sender=celery_worker)

            self.assertFalse(
                Worker.objects.filter(
                    worker_type__in={WorkerType.CELERY, WorkerType.SCHEDULER}
                ).exists()
            )
