# Copyright 2021-2022 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.

"""Tests for scheduler."""

import inspect
import logging
from datetime import timedelta
from typing import TYPE_CHECKING, cast
from unittest import mock

from channels.db import database_sync_to_async

from django.test import override_settings
from django.utils import timezone

from debusine.artifacts.models import ArtifactCategory, CollectionCategory
from debusine.db.models import (
    Artifact,
    Token,
    WorkRequest,
    Worker,
    WorkflowTemplate,
    default_workspace,
)
from debusine.db.tests.utils import RunInParallelTransaction
from debusine.server.scheduler import (
    TaskDatabase,
    _work_request_changed,
    _worker_changed,
    schedule,
    schedule_for_worker,
)
from debusine.server.tasks.noop import ServerNoop
from debusine.tasks.models import LookupMultiple, TaskTypes, WorkerType
from debusine.tasks.sbuild import Sbuild
from debusine.tasks.tests.helper_mixin import FakeTaskDatabase
from debusine.test.django import TestCase, TransactionTestCase

if TYPE_CHECKING:
    from debusine.test.playground import Playground


class HelperMixin:
    """Additional methods to help test scheduler related functions."""

    SAMPLE_TASK_DATA = {
        "input": {
            "source_artifact": 10,
        },
        "distribution": "unstable",
        "backend": "schroot",
        "host_architecture": "amd64",
        "build_components": [
            "any",
            "all",
        ],
    }

    playground: "Playground"

    def patch_task_database(self) -> None:
        """Patch :py:class:`debusine.server.scheduler.TaskDatabase`."""
        task_database_patcher = mock.patch(
            "debusine.server.scheduler.TaskDatabase",
            side_effect=lambda _: FakeTaskDatabase(  # noqa: U101
                single_lookups={(10, None): 10}
            ),
        )
        task_database_patcher.start()
        cast(TestCase, self).addCleanup(task_database_patcher.stop)

    @classmethod
    def _create_worker(
        cls, *, enabled: bool, worker_type: WorkerType = WorkerType.EXTERNAL
    ) -> Worker:
        """Create and return a Worker."""
        token = Token.objects.create(enabled=enabled)
        worker = Worker.objects.create_with_fqdn('worker.lan', token)

        sbuild_chroots = "{}-{}".format(
            cls.SAMPLE_TASK_DATA['distribution'],
            cls.SAMPLE_TASK_DATA['host_architecture'],
        )
        worker.dynamic_metadata = {
            "system:worker_type": worker_type,
            "system:host_architecture": "amd64",
            "system:architectures": ["amd64"],
            "sbuild:version": Sbuild.TASK_VERSION,
            "sbuild:available": True,
            "sbuild:chroots": [sbuild_chroots],
        }

        worker.mark_connected()

        return worker

    def create_sample_sbuild_work_request(
        self, assign_to: Worker | None = None
    ):
        """Create and return a WorkRequest."""
        work_request = self.playground.create_work_request(task_name='sbuild')
        work_request.task_data = dict(self.SAMPLE_TASK_DATA)
        work_request.save()

        if assign_to:
            work_request.assign_worker(assign_to)

        return work_request

    def set_tasks_allowlist(self, allowlist):
        """Set tasks_allowlist in worker's static metadata."""
        self.worker.static_metadata["tasks_allowlist"] = allowlist
        self.worker.save()

    def set_tasks_denylist(self, denylist):
        """Set tasks_denylist in worker's static metadata."""
        self.worker.static_metadata["tasks_denylist"] = denylist
        self.worker.save()


class TaskDatabaseTests(TestCase):
    """Test database interaction in worker tasks."""

    def test_lookup_single_artifact(self) -> None:
        """Look up a single artifact."""
        workspace = self.create_workspace(name="test")
        work_request = self.create_work_request(
            task_name="noop", workspace=workspace
        )
        collection = self.create_collection(
            "debian", CollectionCategory.ENVIRONMENTS, workspace=workspace
        )
        artifact, _ = self.create_artifact(
            category=ArtifactCategory.SYSTEM_TARBALL,
            data={"codename": "bookworm", "architecture": "amd64"},
        )
        collection.manager.add_artifact(artifact, user=self.get_test_user())
        task_db = TaskDatabase(work_request)

        self.assertEqual(
            task_db.lookup_single_artifact(
                "debian@debian:environments/match:codename=bookworm"
            ),
            artifact.id,
        )
        self.assertEqual(
            task_db.lookup_single_artifact(
                "debian/match:codename=bookworm",
                default_category=CollectionCategory.ENVIRONMENTS,
            ),
            artifact.id,
        )
        with self.assertRaisesRegex(
            LookupError,
            "'debian' does not specify a category and the context does not "
            "supply a default",
        ):
            task_db.lookup_single_artifact("debian/match:codename=bookworm")

    def test_lookup_multiple_artifacts(self) -> None:
        """Look up multiple artifacts."""
        workspace = self.create_workspace(name="test")
        work_request = self.create_work_request(
            task_name="noop", workspace=workspace
        )
        collection = self.create_collection(
            "debian", CollectionCategory.ENVIRONMENTS, workspace=workspace
        )
        artifacts: list[Artifact] = []
        for codename in "bookworm", "trixie":
            artifact, _ = self.create_artifact(
                category=ArtifactCategory.SYSTEM_TARBALL,
                data={"codename": codename, "architecture": "amd64"},
            )
            collection.manager.add_artifact(artifact, user=self.get_test_user())
            artifacts.append(artifact)
        task_db = TaskDatabase(work_request)

        self.assertCountEqual(
            task_db.lookup_multiple_artifacts(
                LookupMultiple.parse_obj(
                    {"collection": "debian@debian:environments"}
                )
            ),
            [artifact.id for artifact in artifacts],
        )
        self.assertCountEqual(
            task_db.lookup_multiple_artifacts(
                LookupMultiple.parse_obj({"collection": "debian"}),
                default_category=CollectionCategory.ENVIRONMENTS,
            ),
            [artifact.id for artifact in artifacts],
        )
        with self.assertRaisesRegex(
            LookupError,
            "'debian' does not specify a category and the context does not "
            "supply a default",
        ):
            task_db.lookup_multiple_artifacts(
                LookupMultiple.parse_obj({"collection": "debian"})
            )


class SchedulerTests(HelperMixin, TestCase):
    """Test schedule() function."""

    def setUp(self):
        """Initialize test case."""
        self.work_request_1 = self.create_sample_sbuild_work_request()
        self.work_request_2 = self.create_sample_sbuild_work_request()
        self.patch_task_database()

    def test_no_work_request_available(self):
        """schedule() returns nothing when no work request is available."""
        self._create_worker(enabled=True)
        WorkRequest.objects.all().delete()

        self.assertEqual(schedule(), [])

    def test_only_unhandled_work_requests_available(self):
        """schedule() returns nothing when no work request is available."""
        self._create_worker(enabled=True)
        user = self.get_test_user()
        WorkRequest.objects.all().delete()

        WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=user,
            status=WorkRequest.Statuses.PENDING,
            task_type=TaskTypes.INTERNAL,
            task_name="unhandled",
            task_data={},
        )

        self.assertEqual(schedule(), [])

    def test_synchronization_points(self):
        """schedule() marks pending synchronization points as completed."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()

        parent = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=user,
            status=WorkRequest.Statuses.RUNNING,
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
            task_data={},
        )

        pending = [
            WorkRequest.objects.create_synchronization_point(
                parent=parent, step=f"pending{i}"
            )
            for i in range(2)
        ]
        for wr in pending:
            wr.status = WorkRequest.Statuses.PENDING
            wr.save()
        WorkRequest.objects.create_synchronization_point(
            parent=parent,
            step="completed",
            status=WorkRequest.Statuses.COMPLETED,
        )

        self.assertCountEqual(schedule(), pending)
        for wr in pending:
            wr.refresh_from_db()
            self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
            self.assertEqual(wr.result, WorkRequest.Results.SUCCESS)

    def test_workflow_callbacks(self) -> None:
        """schedule() runs workflow callbacks and marks them as completed."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        parent = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=user,
            status=WorkRequest.Statuses.RUNNING,
            task_type=TaskTypes.WORKFLOW,
            task_name="noop",
            task_data={},
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        wr.status = WorkRequest.Statuses.PENDING
        wr.save()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback"
        ) as mock_noop_callback:
            self.assertEqual(schedule(), [wr])

        mock_noop_callback.assert_called_once_with(wr)
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.SUCCESS)

    def test_workflow_callback_outside_workflow(self) -> None:
        """Workflow callbacks outside workflows are skipped."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=user,
            status=WorkRequest.Statuses.PENDING,
            task_type=TaskTypes.INTERNAL,
            task_name="workflow",
        )

        with self.assertLogs(
            "debusine.server.scheduler", level=logging.WARNING
        ) as log:
            self.assertEqual(schedule(), [])

        self.assertIn(
            f"Workflow callback {wr.id} is not contained in a "
            f"workflow; aborting it.",
            "\n".join(log.output),
        )
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow_callback_bad_task_data(self) -> None:
        """Workflow callbacks with bad task data are skipped."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        parent = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=user,
            status=WorkRequest.Statuses.RUNNING,
            task_type=TaskTypes.WORKFLOW,
            task_name="noop",
            task_data={"nonsense": ""},
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        wr.status = WorkRequest.Statuses.PENDING
        wr.save()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback"
        ) as mock_noop_callback, self.assertLogs(
            "debusine.server.scheduler", level=logging.WARNING
        ) as log:
            self.assertEqual(schedule(), [])

        mock_noop_callback.assert_not_called()
        self.assertIn(
            f"Workflow orchestrator failed to configure for WorkRequest "
            f"{wr.id}; aborting it. Error: 1 validation error",
            "\n".join(log.output),
        )
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow_callback_fails(self) -> None:
        """schedule() logs workflow callbacks that fail."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        parent = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=user,
            status=WorkRequest.Statuses.RUNNING,
            task_type=TaskTypes.WORKFLOW,
            task_name="noop",
            task_data={},
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        wr.status = WorkRequest.Statuses.PENDING
        wr.save()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback",
            side_effect=ValueError("Boom"),
        ) as mock_noop_callback, self.assertLogs(
            "debusine.server.scheduler", level=logging.WARNING
        ) as log:
            self.assertEqual(schedule(), [])

        mock_noop_callback.assert_called_once_with(wr)
        self.assertIn(
            f"Workflow orchestrator failed for WorkRequest {wr.id}. "
            f"Error: Boom",
            "\n".join(log.output),
        )
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflows(self) -> None:
        """schedule() populates workflows and leaves them running."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        template = WorkflowTemplate.objects.create(
            name="test", workspace=default_workspace(), task_name="noop"
        )
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=user
        )

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate"
        ) as mock_noop_populate:
            self.assertEqual(schedule(), [workflow])

        mock_noop_populate.assert_called_once_with()
        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(workflow.result, WorkRequest.Results.NONE)

    def test_workflows_bad_task_data(self) -> None:
        """Workflows with bad task data are skipped."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        template = WorkflowTemplate.objects.create(
            name="test", workspace=default_workspace(), task_name="noop"
        )
        # Bad task data would normally be caught by create_workflow.  Force
        # it to happen here by changing the task data after initial
        # creation.  (A more realistic case might be one where the
        # definition of a workflow changes but existing data isn't
        # migrated.)
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=user
        )
        workflow.task_data = {"nonsense": ""}
        workflow.save()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate"
        ) as mock_noop_populate, self.assertLogs(
            "debusine.server.scheduler", level=logging.WARNING
        ) as log:
            self.assertEqual(schedule(), [])

        mock_noop_populate.assert_not_called()
        self.assertIn(
            f"Workflow orchestrator failed to configure for WorkRequest "
            f"{workflow.id}; aborting it. Error: 1 validation error",
            "\n".join(log.output),
        )
        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(workflow.result, WorkRequest.Results.ERROR)

    def test_workflow_fails(self) -> None:
        """schedule() logs workflows that fail."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        template = WorkflowTemplate.objects.create(
            name="test", workspace=default_workspace(), task_name="noop"
        )
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=user
        )

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate",
            side_effect=ValueError("Boom"),
        ) as mock_noop_populate, self.assertLogs(
            "debusine.server.scheduler", level=logging.WARNING
        ) as log:
            self.assertEqual(schedule(), [])

        mock_noop_populate.assert_called_once_with()
        self.assertIn(
            f"Workflow orchestrator failed for WorkRequest {workflow.id}. "
            f"Error: Boom",
            "\n".join(log.output),
        )
        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(workflow.result, WorkRequest.Results.ERROR)

    def test_workflow_unblocks_children(self) -> None:
        """schedule() unblocks workflow children if possible."""
        user = self.get_test_user()
        WorkRequest.objects.all().delete()
        template = WorkflowTemplate.objects.create(
            name="test", workspace=default_workspace(), task_name="noop"
        )
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=user
        )

        def populate():
            children = [workflow.create_child("noop") for _ in range(2)]
            children[1].add_dependency(children[0])

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate",
            side_effect=populate,
        ) as mock_noop_populate:
            self.assertEqual(schedule(), [workflow])

        mock_noop_populate.assert_called_once_with()
        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(workflow.result, WorkRequest.Results.NONE)
        children = workflow.children.order_by("id")
        self.assertEqual(children.count(), 2)
        self.assertEqual(children[0].status, WorkRequest.Statuses.PENDING)
        self.assertEqual(children[1].status, WorkRequest.Statuses.BLOCKED)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_worker_connected_without_a_work_request(self):
        """schedule() goes over all workers."""
        self._create_worker(enabled=True)
        self._create_worker(enabled=True)

        result = schedule()

        self.assertCountEqual(
            result, {self.work_request_1, self.work_request_2}
        )

    def test_worker_connected_with_a_work_request(self):
        """schedule() skips a worker with assigned work requests."""
        # Create worker_1 and assigns to work_request_1
        worker_1 = self._create_worker(enabled=True)
        self.work_request_1.assign_worker(worker_1)

        # schedule() has no worker available
        self.assertEqual(schedule(), [])

    def test_worker_not_connected(self):
        """schedule() skips the workers that are not connected."""
        worker = self._create_worker(enabled=True)
        worker.connected_at = None
        worker.save()

        self.assertEqual(schedule(), [])

    def test_worker_not_enabled(self):
        """schedule() skips workers that are not enabled."""
        worker = self._create_worker(enabled=False)
        worker.mark_connected()

        self.assertEqual(schedule(), [])


class ScheduleForWorkerTests(HelperMixin, TestCase):
    """Test schedule_for_worker()."""

    def setUp(self) -> None:
        """Initialize test case."""
        self.worker = self._create_worker(enabled=True)
        self.patch_task_database()

    def test_no_work_request(self):
        """schedule_for_worker() returns None when nothing to do."""
        self.assertIsNone(schedule_for_worker(self.worker))

    def test_skips_assigned_work_requests(self):
        """schedule_for_worker() doesn't pick assigned work requests."""
        # Create a running work request
        worker_2 = self._create_worker(enabled=True)
        work_request_1 = self.create_sample_sbuild_work_request()
        work_request_1.assign_worker(worker_2)

        result = schedule_for_worker(self.worker)

        self.assertIsNone(result)

    def test_skips_non_pending_work_requests(self):
        """schedule_for_worker() skips non-pending work requests."""
        # Create a completed work request
        worker_2 = self._create_worker(enabled=True)
        work_request_1 = self.create_sample_sbuild_work_request(
            assign_to=worker_2
        )
        work_request_1.mark_running()
        work_request_1.mark_completed(WorkRequest.Results.SUCCESS)
        # Create a running work request
        work_request_2 = self.create_sample_sbuild_work_request(
            assign_to=worker_2
        )
        work_request_2.mark_running()
        # Create an aborted work request
        work_request_2 = self.create_sample_sbuild_work_request()
        work_request_2.mark_aborted()

        result = schedule_for_worker(self.worker)

        self.assertIsNone(result)

    def test_skips_work_request_based_on_can_run_on(self):
        """schedule_for_worker() calls task->can_run_on."""
        self.create_sample_sbuild_work_request()
        # The request is for unstable-amd64, so can_run_on will return False
        self.worker.dynamic_metadata = {
            "system.worker_type": WorkerType.EXTERNAL,
            "sbuild:chroots": ["stable-amd64"],
        }
        self.worker.save()

        result = schedule_for_worker(self.worker)

        self.assertIsNone(result)

    def test_with_pending_unassigned_work_request(self):
        """schedule_for_worker() picks up the pending work request."""
        with override_settings(DISABLE_AUTOMATIC_SCHEDULING=True):
            work_request = self.create_sample_sbuild_work_request()

        with self.captureOnCommitCallbacks() as on_commit:
            result = schedule_for_worker(self.worker)

        self.assertIsInstance(result, WorkRequest)
        self.assertEqual(result, work_request)
        self.assertEqual(result.worker, self.worker)

        # Assigning the work request does not cause another invocation of the
        # scheduler.
        self.assertEqual(on_commit, [])

    def test_with_failing_configure_call(self):
        """schedule_for_worker() skips invalid data and marks it as an error."""
        work_request = self.create_sample_sbuild_work_request()
        work_request.task_data = {}  # invalid: it's missing required properties
        work_request.save()

        with self.assertLogs(
            "debusine.server.scheduler", level=logging.WARNING
        ) as log:
            result = schedule_for_worker(self.worker)
        work_request.refresh_from_db()

        self.assertIn(
            f"WorkRequest {work_request.id} failed to configure, aborting it. "
            f"Task data: {work_request.task_data} Error: 3 validation errors",
            "\n".join(log.output),
        )
        self.assertIsNone(result)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.ERROR)

    def test_with_failing_compute_dynamic_data(self):
        """schedule_for_worker() handles failure to compute dynamic data."""
        work_request = self.create_sample_sbuild_work_request()
        work_request.task_data["backend"] = "unshare"
        del work_request.task_data["distribution"]
        work_request.task_data["environment"] = "bad-lookup"
        work_request.save()
        self.worker.dynamic_metadata["executor:unshare:available"] = True

        with self.assertLogs(
            "debusine.server.scheduler", level=logging.WARNING
        ) as log:
            result = schedule_for_worker(self.worker)
        work_request.refresh_from_db()

        self.assertIn(
            f"WorkRequest {work_request.id} failed to pre-process task data, "
            f"aborting it. "
            f"Task data: {work_request.task_data} "
            f"Error: ('bad-lookup', {CollectionCategory.ENVIRONMENTS!r})",
            "\n".join(log.output),
        )
        self.assertIsNone(result)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.ERROR)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_skips_tasks_in_deny_list(self):
        """schedule_for_worker() doesn't schedule task in denylist."""
        self.create_sample_sbuild_work_request()  # is sbuild task
        self.set_tasks_denylist(["foobar", "sbuild", "baz"])

        self.assertIsNone(schedule_for_worker(self.worker))

    def test_skips_tasks_not_in_allow_list(self):
        """schedule_for_worker() doesn't schedule task in denylist."""
        self.create_sample_sbuild_work_request()  # is sbuild task
        self.set_tasks_allowlist(["foobar", "baz"])

        self.assertIsNone(schedule_for_worker(self.worker))

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_allowlist_takes_precedence_over_denylist(self):
        """A task that is in both allow/denylist is allowed."""
        work_request = (
            self.create_sample_sbuild_work_request()
        )  # is sbuild task
        self.set_tasks_allowlist(["sbuild"])
        self.set_tasks_denylist(["sbuild"])

        result = schedule_for_worker(self.worker)

        self.assertEqual(work_request, result)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_highest_priority_work_request_is_picked_first(self):
        """schedule_for_worker() picks the highest-priority work request."""
        list_of_work_requests = []
        for priority_base, priority_adjustment in (
            (0, 0),
            (0, 10),
            (10, 0),
            (10, 1),
        ):
            work_request = self.create_sample_sbuild_work_request()
            work_request.priority_base = priority_base
            work_request.priority_adjustment = priority_adjustment
            work_request.save()
            list_of_work_requests.append(work_request)

        result = schedule_for_worker(self.worker)

        self.assertEqual(list_of_work_requests[3], result)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_oldest_work_request_is_picked_first(self):
        """schedule_for_worker() picks the oldest work request."""
        now = timezone.now()
        list_of_work_requests = []
        for offset in (1, 0, 2):  # The second object created will be the oldest
            work_request = self.create_sample_sbuild_work_request()
            work_request.created_at = now + timedelta(seconds=offset)
            work_request.save()
            list_of_work_requests.append(work_request)

        result = schedule_for_worker(self.worker)

        self.assertEqual(list_of_work_requests[1], result)

    def test_signing_work_request(self):
        """schedule_for_worker() assigns a signing work request."""
        signing_worker = self._create_worker(
            enabled=True, worker_type=WorkerType.SIGNING
        )
        with override_settings(DISABLE_AUTOMATIC_SCHEDULING=True):
            work_request = self.playground.create_work_request(
                task_type=TaskTypes.SIGNING,
                task_name="generatekey",
                task_data={"purpose": "uefi", "description": "A UEFI key"},
            )

        result = schedule_for_worker(signing_worker)

        self.assertEqual(result, work_request)
        self.assertEqual(result.worker, signing_worker)

    def assert_schedule_for_worker_return_none(self, status):
        """Schedule_for_worker returns None if it has WorkRequest in status."""
        work_request = self.create_sample_sbuild_work_request(
            assign_to=self.worker
        )
        self.create_sample_sbuild_work_request()

        work_request.status = status
        work_request.save()

        self.assertIsNone(schedule_for_worker(self.worker))

    def test_schedule_for_worker_already_pending(self):
        """Do not assign a WorkRequest if the Worker has enough pending."""
        self.assert_schedule_for_worker_return_none(
            WorkRequest.Statuses.PENDING
        )

    def test_schedule_for_worker_already_running(self):
        """Do not assign a WorkRequest if the Worker has enough running."""
        self.assert_schedule_for_worker_return_none(
            WorkRequest.Statuses.RUNNING
        )

    @mock.patch("debusine.server.scheduler.run_server_task.apply_async")
    def test_schedule_for_worker_concurrency(
        self, mock_apply_async  # noqa: U100
    ):
        """Consider a Worker's concurrency when assigning a WorkRequest."""
        worker = Worker.objects.get_or_create_celery()
        worker.concurrency = 3
        worker.dynamic_metadata = {
            "system:worker_type": WorkerType.CELERY,
            "server:servernoop:version": ServerNoop.TASK_VERSION,
        }

        for _ in range(3):
            work_request = self.playground.create_work_request(
                task_type=TaskTypes.SERVER, task_name="servernoop"
            )
            self.assertEqual(work_request, schedule_for_worker(worker))

        work_request = self.playground.create_work_request(
            task_type=TaskTypes.SERVER, task_name="servernoop"
        )
        self.assertIsNone(schedule_for_worker(worker))

    def test_work_request_changed_has_kwargs(self):
        """
        Check method _work_request_changed has kwargs argument.

        If it doesn't have the signal connection fails if DEBUG=0
        """
        _method_has_kwargs(_work_request_changed)

    def test_worker_changed_has_kwargs(self):
        """
        Check method _worker_changed has kwargs argument.

        If it doesn't have the signal connection fails if DEBUG=0
        """
        _method_has_kwargs(_worker_changed)


class ScheduleForWorkerTransactionTests(HelperMixin, TransactionTestCase):
    """Test schedule_for_worker()."""

    def setUp(self):
        """Initialize test case."""
        super().setUp()
        self.worker = self._create_worker(enabled=True)
        self.patch_task_database()

    @override_settings(
        DISABLE_AUTOMATIC_SCHEDULING=False, CELERY_TASK_ALWAYS_EAGER=True
    )
    async def test_schedule_from_work_request_changed(self):
        """End to end: from WorkRequest.save(), notification to schedule()."""
        work_request_1 = await database_sync_to_async(
            self.create_sample_sbuild_work_request
        )()

        self.assertIsNone(work_request_1.worker)

        # When work_request_1 got created the function
        # _work_request_changed() was called (because it is connected
        # to the signal post_save on the WorkRequest). It called then
        # schedule() which assigned the work_request_1 to worker
        await database_sync_to_async(work_request_1.refresh_from_db)()
        work_request_1_worker = await database_sync_to_async(
            lambda: work_request_1.worker
        )()

        self.assertEqual(work_request_1.worker, work_request_1_worker)

        # Create another work request
        work_request_2 = await database_sync_to_async(
            self.create_sample_sbuild_work_request
        )()
        self.assertIsNone(work_request_2.worker)

        await database_sync_to_async(work_request_2.refresh_from_db)()
        work_request_2_worker = await database_sync_to_async(
            lambda: work_request_2.worker
        )()

        # schedule() got called but there isn't any available worker
        # (work_request_1 is assigned to self.worker and has not finished)
        self.assertIsNone(work_request_2_worker)

        # Mark work_request_1 as done transitioning to RUNNING, then
        # mark_completed(). schedule() got called when work_request_1 had
        # changed and then it assigned work_request_2 to self.worker
        self.assertTrue(
            await database_sync_to_async(work_request_1.mark_running)()
        )
        await database_sync_to_async(work_request_1.mark_completed)(
            WorkRequest.Results.SUCCESS
        )

        await database_sync_to_async(work_request_2.refresh_from_db)()
        work_request_2_worker = await database_sync_to_async(
            lambda: work_request_2.worker
        )()

        # Assert that work_request_2 got assigned to self.worker
        self.assertEqual(work_request_2_worker, self.worker)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=False)
    async def test_schedule_from_worker_changed(self):
        """End to end: from Worker.save() to schedule_for_worker()."""
        await database_sync_to_async(self.worker.mark_disconnected)()

        work_request_1 = await database_sync_to_async(
            self.create_sample_sbuild_work_request
        )()

        # The Worker is not connected: the work_request_1.worker is None
        self.assertIsNone(work_request_1.worker)

        # _worker_changed() will be executed via the post_save signal
        # from Worker. It will call schedule_for_worker()
        await database_sync_to_async(self.worker.mark_connected)()

        await database_sync_to_async(work_request_1.refresh_from_db)()

        work_request_1_worker = await database_sync_to_async(
            lambda: work_request_1.worker
        )()

        self.assertEqual(work_request_1_worker, self.worker)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_when_worker_is_locked(self):
        """schedule_for_worker() fails when it fails to lock the Worker."""
        self.create_sample_sbuild_work_request()

        thread = RunInParallelTransaction(
            lambda: Worker.objects.select_for_update().get(id=self.worker.id)
        )
        thread.start_transaction()

        try:
            self.assertIsNone(schedule_for_worker(self.worker))
        finally:
            thread.stop_transaction()

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_when_work_request_is_locked(self):
        """schedule_for_worker() fails when it fails to lock the WorkRequest."""
        work_request = self.create_sample_sbuild_work_request()

        thread = RunInParallelTransaction(
            lambda: WorkRequest.objects.select_for_update().get(
                id=work_request.id
            )
        )
        thread.start_transaction()

        try:
            self.assertIsNone(schedule_for_worker(self.worker))
        finally:
            thread.stop_transaction()

    @mock.patch("debusine.server.scheduler.run_server_task.apply_async")
    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_server_task(self, mock_apply_async):
        """schedule_for_worker() dispatches server tasks to Celery."""
        worker = Worker.objects.get_or_create_celery()
        worker.dynamic_metadata = {
            "system:worker_type": WorkerType.CELERY,
            "server:servernoop:version": ServerNoop.TASK_VERSION,
        }
        work_request = self.playground.create_work_request(
            task_type=TaskTypes.SERVER, task_name="servernoop"
        )

        result = schedule_for_worker(worker)

        self.assertEqual(work_request, result)
        mock_apply_async.assert_called_once()
        self.assertEqual(
            mock_apply_async.call_args.kwargs["args"], (work_request.id,)
        )
        self.assertRegex(
            mock_apply_async.call_args.kwargs["task_id"],
            fr"^servernoop_{work_request.id}_"
            r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",
        )


def _method_has_kwargs(method):
    """
    Assert method has kwargs: suitable for being a signal.

    This is enforced by Django when settings.DEBUG=True. Tests are
    executed with settings.DEBUG=False so the test re-implements Django
    validation to make sure that `**kwargs` is a parameter in the method
    (otherwise is detected only in production).
    """
    parameters = inspect.signature(method).parameters

    return any(  # pragma: no cover
        p for p in parameters.values() if p.kind == p.VAR_KEYWORD
    )
