# 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.

"""Tests for scheduler."""

import datetime as dt
import logging
from contextlib import ExitStack
from typing import Any, TYPE_CHECKING, override
from unittest import mock

from celery import states
from celery.contrib.testing.app import TestApp, setup_default_app
from channels.db import database_sync_to_async
from django.db import OperationalError, transaction
from django.db.models import QuerySet
from django.test import override_settings
from django.utils import timezone
from django_celery_results.models import TaskResult

import debusine.tasks.tags as ttags
import debusine.worker.tags as wtags
from debusine.artifacts.models import CollectionCategory
from debusine.assets import KeyPurpose
from debusine.db.context import context
from debusine.db.models import CollectionItem, WorkRequest, Worker
from debusine.db.models.lookups import LookupCache
from debusine.db.models.work_requests import workflow_ancestors
from debusine.db.playground import scenarios
from debusine.db.tasks.mock_resolver import MockFieldResolver
from debusine.db.tests.utils import RunInParallelTransaction
from debusine.project.celery import (
    TaskWithBetterArgumentRepresentation,
    celery_tasks_in_progress,
    get_apply_async_tasks,
)
from debusine.server.scheduler import (
    get_provided_worker_tags,
    schedule,
    schedule_for_worker,
    schedule_soon,
    schedule_task,
)
from debusine.server.tasks import ServerNoop
from debusine.server.tasks.tests.helpers import SampleBaseWaitTask
from debusine.server.tasks.wait.models import ConfirmData, DelayData
from debusine.server.workflows.models import WorkRequestWorkflowData
from debusine.signing.tasks.models import GenerateKeyData
from debusine.tasks import Lintian, Sbuild
from debusine.tasks.models import (
    ActionSkipIfLookupResultChanged,
    BackendType,
    BaseDynamicTaskData,
    BaseTaskData,
    Confirmation,
    EventReactions,
    LintianData,
    LintianInput,
    OutputData,
    OutputDataError,
    SbuildBuildComponent,
    SbuildData,
    SbuildInput,
    WorkerType,
)
from debusine.test.django import (
    PlaygroundTestCase,
    TestCase,
    TransactionTestCase,
    preserve_db_task_registry,
)
from debusine.worker.system_information import native_architecture

if TYPE_CHECKING:
    from debusine.test.playground import Playground


class SchedulerTestCase(PlaygroundTestCase):
    """Additional methods to help test scheduler related functions."""

    scenario = scenarios.DefaultSbuild()

    SAMPLE_TASK_DATA = SbuildData(
        input=SbuildInput(source_artifact=10),
        environment="debian/match:codename=sid",
        build_architecture="amd64",
        build_components=[
            SbuildBuildComponent("any"),
            SbuildBuildComponent.ANY,
            SbuildBuildComponent.ALL,
        ],
    )

    playground: "Playground"

    def setUp(self) -> None:
        super().setUp()
        self.task_data = self.SAMPLE_TASK_DATA.model_copy()
        self.task_data.input.source_artifact = self.scenario.source.id

        self.resolver = MockFieldResolver(
            extra_repositories=None,
            debusine_fqdn="debusine.example.net",
            task_configuration=None,
        )
        self.resolver.set_artifact("environment", self.scenario.environment_sid)
        self.resolver.set_artifact(
            "input_source_artifact", self.scenario.source
        )
        self.resolver.set_field_artifacts("input_binary_artifacts")
        self.resolver.set_field_artifacts("input_extra_binary_artifacts")
        self.resolver.set_collection(
            "task_configuration",
            self.scenario.default_task_configuration_collection,
        )

    def patch_field_resolver(self) -> None:
        """Patch :py:class:`debusine.server.scheduler.TaskDatabase`."""
        field_resolver_patcher = mock.patch(
            "debusine.db.tasks.base.DBFieldResolver",
            side_effect=lambda _: self.resolver,
        )
        field_resolver_patcher.start()
        self.addCleanup(field_resolver_patcher.stop)

    @transaction.atomic
    def _create_worker(
        self, *, enabled: bool, worker_type: WorkerType = WorkerType.EXTERNAL
    ) -> Worker:
        """Create and return a Worker."""
        token = self.playground.create_worker_token(
            enabled=enabled, set_worker=False
        )
        worker = Worker.objects.create_with_fqdn(
            'worker.lan', token=token, worker_type=worker_type
        )

        worker.mark_connected()

        worker.set_dynamic_metadata(
            {
                "system:worker_type": worker_type,
                "system:native_architecture": "amd64",
                "system:architectures": ["amd64"],
                "sbuild:version": Sbuild.TASK_VERSION,
                "sbuild:available": True,
                "signing:generatekey:version": 1,
                "executor:unshare:available": True,
            }
        )

        return worker

    def create_sample_sbuild_work_request(
        self,
        require_scheduler_tags: bool = True,
        assign_to: Worker | None = None,
        parent: WorkRequest | None = None,
    ) -> WorkRequest:
        """Create and return a WorkRequest."""
        work_request = self.playground.create_worker_task(
            task_name="sbuild", task_data=self.SAMPLE_TASK_DATA, parent=parent
        )
        with ExitStack() as stack:
            if not require_scheduler_tags:
                stack.enter_context(
                    mock.patch(
                        "debusine.db.tasks.base.DBProxyTask"
                        ".compute_scheduler_tags",
                        return_value=(
                            ttags.TaskProvidedTags(),
                            wtags.TaskRequiredTags(),
                        ),
                    )
                )
            self.playground.advance_work_request(
                work_request, mark_pending=True, worker=assign_to
            )
        return work_request

    def set_tasks_allowlist(self, allowlist: list[str]) -> None:
        """Set tasks_allowlist in worker's static metadata."""
        assert hasattr(self, "worker")
        self.worker.static_metadata["tasks_allowlist"] = allowlist
        self.worker.save()

    def set_tasks_denylist(self, denylist: list[str]) -> None:
        """Set tasks_denylist in worker's static metadata."""
        assert hasattr(self, "worker")
        self.worker.static_metadata["tasks_denylist"] = denylist
        self.worker.save()

    def create_wait_delay_task(
        self,
        delay_offset: dt.timedelta,
    ) -> WorkRequest:
        """Create a Wait/Delay task in a workflow."""
        parent = self.playground.create_workflow(mark_running=True)
        return parent.create_child_wait(
            task_name="delay",
            task_data=DelayData(delay_until=timezone.now() + delay_offset),
            workflow_data=WorkRequestWorkflowData(needs_input=False),
        )


class SchedulerTests(SchedulerTestCase, TransactionTestCase):
    """Test schedule() function."""

    def setUp(self) -> None:
        """Initialize test case."""
        super().setUp()
        self.work_request_1 = self.create_sample_sbuild_work_request()
        self.work_request_2 = self.create_sample_sbuild_work_request()
        self.patch_field_resolver()

    def test_no_work_request_available(self) -> None:
        """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) -> None:
        """schedule() returns nothing when no work request is available."""
        self._create_worker(enabled=True)
        WorkRequest.objects.all().delete()

        self.playground.create_internal_task(
            mark_pending=True, task_name="unhandled", validate=False
        )

        self.assertEqual(schedule(), [])

    def test_synchronization_points(self) -> None:
        """schedule() marks pending synchronization points as completed."""
        WorkRequest.objects.all().delete()

        parent = self.playground.create_workflow(mark_running=True)

        pending = [
            WorkRequest.objects.create_synchronization_point(
                parent=parent, step=f"pending{i}"
            )
            for i in range(2)
        ]
        for wr in pending:
            wr.mark_pending()
        self.playground.advance_work_request(
            WorkRequest.objects.create_synchronization_point(
                parent=parent,
                step="completed",
            ),
            result=WorkRequest.Results.SUCCESS,
        )

        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)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_worker_connected_without_a_work_request(self) -> None:
        """schedule() goes over all workers."""
        self.work_request_1.scheduler_tags_required = []
        self.work_request_1.save()
        self.work_request_2.scheduler_tags_required = []
        self.work_request_2.save()
        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) -> None:
        """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) -> None:
        """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) -> None:
        """schedule() skips workers that are not enabled."""
        worker = self._create_worker(enabled=False)
        worker.mark_connected()

        self.assertEqual(schedule(), [])

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_wait_delay_task_past(self) -> None:
        """schedule() marks expired Wait/Delay tasks as completed."""
        work_request = self.create_wait_delay_task(dt.timedelta(seconds=-1))
        work_request.mark_pending()

        result = schedule()

        self.assertEqual(result, [work_request])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.SUCCESS)
        self.assertIsNone(work_request.output_data)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_wait_delay_task_future(self) -> None:
        """schedule() leaves non-expired Wait/Delay tasks running."""
        work_request = self.create_wait_delay_task(dt.timedelta(hours=1))
        work_request.mark_pending()

        result = schedule()

        self.assertEqual(result, [work_request])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_wait_delay_task_blocked(self) -> None:
        """schedule() skips blocked Wait/Delay tasks."""
        work_request = self.create_wait_delay_task(dt.timedelta(seconds=-1))

        result = schedule()

        self.assertEqual(result, [])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.BLOCKED)

    def create_wait_confirm_task(
        self, *, auto_confirm_if_no_failure: bool
    ) -> WorkRequest:
        """Create a Wait/Confirm task in a workflow."""
        parent = self.playground.create_workflow(mark_running=True)
        return parent.create_child_wait(
            task_name="confirm",
            task_data=ConfirmData(
                auto_confirm_if_no_failure=auto_confirm_if_no_failure
            ),
            workflow_data=WorkRequestWorkflowData(needs_input=True),
        )

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_wait_confirm_task_auto_confirm_dep_success(self) -> None:
        """schedule() handles auto-confirming tasks where all deps succeeded."""
        work_request = self.create_wait_confirm_task(
            auto_confirm_if_no_failure=True
        )
        assert work_request.parent is not None
        dep = work_request.parent.create_child_worker(task_name="noop")
        work_request.add_dependency(dep)
        self.playground.advance_work_request(
            dep, result=WorkRequest.Results.SUCCESS
        )

        result = schedule()

        self.assertEqual(result, [work_request])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.SUCCESS)
        self.assertEqual(
            work_request.output_data,
            OutputData(confirmation=Confirmation(confirmed=True)),
        )

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_wait_confirm_task_auto_confirm_dep_failure(self) -> None:
        """schedule() skips auto-confirming tasks where a dep failed."""
        work_request = self.create_wait_confirm_task(
            auto_confirm_if_no_failure=True
        )
        assert work_request.parent is not None
        dep = work_request.parent.create_child_worker(
            task_name="noop",
            workflow_data=WorkRequestWorkflowData(allow_failure=True),
        )
        work_request.add_dependency(dep)
        self.playground.advance_work_request(
            dep, result=WorkRequest.Results.FAILURE
        )

        result = schedule()

        self.assertEqual(result, [work_request])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_wait_confirm_task_non_auto_confirm(self) -> None:
        """schedule() skips non-auto-confirming Confirm tasks."""
        work_request = self.create_wait_confirm_task(
            auto_confirm_if_no_failure=False
        )
        work_request.mark_pending()

        result = schedule()

        self.assertEqual(result, [work_request])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_wait_confirm_task_blocked(self) -> None:
        """schedule() skips blocked Wait/Confirm tasks."""
        work_request = self.create_wait_confirm_task(
            auto_confirm_if_no_failure=True
        )

        result = schedule()

        self.assertEqual(result, [])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.BLOCKED)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    @preserve_db_task_registry()
    def test_wait_other_task(self) -> None:
        """schedule() leaves non-Delay Wait tasks running."""

        class WaitNoop(SampleBaseWaitTask[BaseTaskData, BaseDynamicTaskData]):
            TASK_VERSION = 1

            @override
            def _execute(self) -> WorkRequest.Results:
                raise NotImplementedError()

        parent = self.playground.create_workflow(mark_running=True)
        work_request = parent.create_child_wait(
            task_name="waitnoop",
            workflow_data=WorkRequestWorkflowData(needs_input=False),
        )
        work_request.mark_pending()

        result = schedule()

        self.assertEqual(result, [work_request])
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)

    def test_schedule_internal_locks_workflows(self) -> None:
        with transaction.atomic():
            work_request = self.create_wait_delay_task(dt.timedelta(hours=1))
            self.playground.advance_work_request(
                work_request, mark_pending=True
            )
            assert work_request.parent is not None
            noop = work_request.parent.create_child_worker(task_name="noop")
            unrelated = self.playground.create_worker_task()

        with (
            transaction.atomic(),
            mock.patch("debusine.server.scheduler.update_workflows"),
        ):
            schedule()

            work_request.refresh_from_db()
            self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)
            self.assert_work_requests_locked(
                locked=[work_request.parent, work_request, noop],
                unlocked=[unrelated],
            )

    def test_schedule_internal_workflow_already_locked(self) -> None:
        with transaction.atomic():
            work_request = self.create_wait_delay_task(dt.timedelta(hours=1))
            self.playground.advance_work_request(
                work_request, mark_pending=True
            )
            assert work_request.parent is not None
            noop = work_request.parent.create_child_worker(task_name="noop")

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

        try:
            with (
                mock.patch("debusine.server.scheduler.update_workflows"),
                self.assertLogsContains(
                    "[SCHEDULER] Failed to lock workflows with pending internal"
                    " work requests",
                    logger="debusine.server.scheduler",
                    level=logging.DEBUG,
                ),
            ):
                schedule()
        finally:
            thread.stop_transaction()

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_complete_wait_tasks_locks_workflows(self) -> None:
        with transaction.atomic():
            work_request = self.create_wait_delay_task(dt.timedelta(seconds=-1))
            self.playground.advance_work_request(
                work_request, mark_running=True
            )
            assert work_request.parent is not None
            noop = work_request.parent.create_child_worker(task_name="noop")
            unrelated = self.playground.create_worker_task()

        with (
            transaction.atomic(),
            mock.patch("debusine.server.scheduler.update_workflows"),
        ):
            schedule()

            work_request.refresh_from_db()
            self.assertEqual(
                work_request.status, WorkRequest.Statuses.COMPLETED
            )
            self.assertEqual(work_request.result, WorkRequest.Results.SUCCESS)
            self.assert_work_requests_locked(
                locked=[work_request.parent, work_request, noop],
                unlocked=[unrelated],
            )

    def test_complete_wait_tasks_workflow_already_locked(self) -> None:
        with transaction.atomic():
            work_request = self.create_wait_delay_task(dt.timedelta(seconds=-1))
            self.playground.advance_work_request(
                work_request, mark_running=True
            )
            assert work_request.parent is not None
            noop = work_request.parent.create_child_worker(task_name="noop")

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

        try:
            with (
                mock.patch("debusine.server.scheduler.update_workflows"),
                self.assertLogsContains(
                    "[SCHEDULER] Failed to lock workflows with pending internal"
                    " work requests",
                    logger="debusine.server.scheduler",
                    level=logging.DEBUG,
                ),
            ):
                schedule()
        finally:
            thread.stop_transaction()


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

    def setUp(self) -> None:
        """Initialize test case."""
        super().setUp()
        self.worker = self._create_worker(enabled=True)
        self.patch_field_resolver()
        self.enterContext(override_settings(FEATURE_TAG_BASED_SCHEDULING=False))

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

    def test_skips_assigned_work_requests(self) -> None:
        """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 = self.playground.schedule_for_worker(self.worker)

        self.assertIsNone(result)

    def test_skips_non_pending_work_requests(self) -> None:
        """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 = self.playground.schedule_for_worker(self.worker)

        self.assertIsNone(result)

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_skips_work_request_based_on_worker_metadata(self) -> None:
        wr = self.create_sample_sbuild_work_request()
        self.worker.set_dynamic_metadata({})
        self.worker.save()
        wr.refresh_from_db()
        self.assertIsNone(wr.worker)

        result = self.playground.schedule_for_worker(self.worker)
        self.assertIsNone(result)

    def test_skips_work_request_if_lookup_result_changed(self) -> None:
        """Skip work request if an action requests it."""
        collection = self.playground.create_collection(
            "test", CollectionCategory.TEST
        )
        lookup = f"test@{CollectionCategory.TEST}/name:foo"
        work_request = self.create_sample_sbuild_work_request(
            require_scheduler_tags=False
        )
        work_request.event_reactions = EventReactions(
            on_assignment=[
                ActionSkipIfLookupResultChanged(
                    lookup=lookup, collection_item_id=None
                )
            ]
        )
        work_request.save()
        CollectionItem.objects.create_from_artifact(
            self.playground.create_artifact()[0],
            parent_collection=collection,
            name="foo",
            data={},
            created_by_user=self.playground.get_default_user(),
        )

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

        self.assertIn(
            f"DEBUG:debusine.server.scheduler:"
            f"Skipping work request {work_request} due to on_assignment action",
            log.output,
        )
        self.assertIsNone(result)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.SKIPPED)
        self.assertEqual(
            work_request.output_data,
            OutputData(skip_reason=f"Result of lookup {lookup!r} changed"),
        )

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

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

        assert isinstance(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, [])

        # Lookups from computing dynamic data are not cached.
        self.assertEqual(context.lookup_cache, LookupCache())

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

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

        expected_message = (
            "Failed to configure: 3 validation errors for SbuildData"
        )
        self.assertIn(
            f"ERROR:debusine.server.scheduler:"
            f"Error assigning work request "
            f"{work_request.task_type}/{work_request.task_name} "
            f"({work_request.id}): {expected_message}",
            "\n".join(log.output),
        )
        self.assertIsNone(result)
        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_with_failing_compute_dynamic_data(self) -> None:
        """schedule_for_worker() handles failure to compute dynamic data."""
        work_request = self.create_sample_sbuild_work_request(
            require_scheduler_tags=False
        )
        work_request.task_data["backend"] = BackendType.UNSHARE
        work_request.task_data["environment"] = "bad-lookup"
        work_request.save()
        self.worker.dynamic_metadata["executor:unshare:available"] = True

        with (
            mock.patch(
                "debusine.server.scheduler.WorkRequest.assign_worker",
                side_effect=LookupError("expected failure"),
            ),
            self.assertLogs(
                "debusine.server.scheduler", level=logging.WARNING
            ) as log,
        ):
            result = schedule_for_worker(self.worker, set())
        work_request.refresh_from_db()

        expected_message = "Failed to compute dynamic data: expected failure"
        self.assertEqual(
            log.output,
            [
                f"ERROR:debusine.server.scheduler:"
                f"Error assigning work request "
                f"{work_request.task_type}/{work_request.task_name} "
                f"({work_request.id}): {expected_message}"
            ],
        )

        self.assertIsNone(result)
        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="dynamic-data-failed"
                    )
                ]
            ),
        )

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_computes_dynamic_data(
        self,
    ) -> None:
        """schedule_for_worker() uses worker arch to compute dynamic data."""
        self.patch_field_resolver()
        self.worker.dynamic_metadata.update(
            {
                "system:native_architecture": "i386",
                "system:architectures": ["amd64", "i386"],
                "lintian:version": Lintian.TASK_VERSION,
                "provided_tags": [
                    "worker:executor:unshare",
                    "worker:task:worker:lintian:version:1",
                    "worker:type:external",
                    wtags.BUILD_ARCH_PREFIX + "amd64",
                    wtags.BUILD_ARCH_PREFIX + "i386",
                ],
            }
        )
        # Lintian tasks don't care what architecture they run on, so they
        # defer to the worker's native architecture.
        work_request = self.playground.create_worker_task(
            task_name="lintian",
            task_data=LintianData(
                input=LintianInput(source_artifact=1),
                environment="debian/match:codename=bookworm",
            ),
            mark_pending=True,
        )
        work_request.save()

        result = self.playground.schedule_for_worker(self.worker)
        assert isinstance(result, WorkRequest)
        self.assertEqual(result, work_request)
        self.assertEqual(result.worker, self.worker)

        task_config_collection = (
            self.scenario.default_task_configuration_collection
        )

        self.assertEqual(
            result.configured_task_data,
            work_request.task_data
            | {"task_configuration": task_config_collection.pk},
        )
        assert result.dynamic_task_data is not None
        dynamic_data = result.dynamic_task_data.copy()
        task_input_fields = dynamic_data.pop("task_input_fields")
        self.assertEqual(
            dynamic_data,
            {
                "input_source_artifact_id": self.scenario.source.id,
                "input_binary_artifacts_ids": [],
                "environment_id": self.scenario.environment_sid.id,
                "subject": "hello",
                "parameter_summary": "hello_1.0-1",
                "runtime_context": "binary-all+binary-any+source:sid",
                "configuration_context": "sid",
            },
        )
        self.assertCountEqual(
            task_input_fields.keys(),
            [
                "build_architecture",
                "environment",
                "input_binary_artifacts",
                "input_source_artifact",
                "task_configuration",
            ],
        )

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_skips_tasks_in_deny_list(self) -> None:
        """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(self.playground.schedule_for_worker(self.worker))

    def test_skips_tasks_not_in_allow_list(self) -> None:
        """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(self.playground.schedule_for_worker(self.worker))

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

        result = self.playground.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) -> None:
        """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(
                require_scheduler_tags=False
            )
            work_request.priority_base = priority_base
            work_request.priority_adjustment = priority_adjustment
            work_request.save()
            list_of_work_requests.append(work_request)

        result = self.playground.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) -> None:
        """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(
                require_scheduler_tags=False
            )
            work_request.created_at = now + dt.timedelta(seconds=offset)
            work_request.save()
            list_of_work_requests.append(work_request)

        result = self.playground.schedule_for_worker(self.worker)

        self.assertEqual(list_of_work_requests[1], result)

    def test_signing_work_request(self) -> None:
        """schedule_for_worker() assigns a signing work request."""
        signing_worker = self._create_worker(
            enabled=True,
            worker_type=WorkerType.SIGNING,
        )
        signing_worker.static_metadata = {
            "provided_tags": [
                wtags.WORKER_TYPE_SIGNING,
                wtags.TASK_PREFIX + "signing:generatekey:version:1",
                wtags.BUILD_ARCH_PREFIX + "amd64",
            ]
        }
        signing_worker.save()
        with override_settings(DISABLE_AUTOMATIC_SCHEDULING=True):
            work_request = self.playground.create_signing_task(
                task_name="generatekey",
                task_data=GenerateKeyData(
                    purpose=KeyPurpose.UEFI, description="A UEFI key"
                ),
                mark_pending=True,
            )

        result = self.playground.schedule_for_worker(signing_worker)

        assert isinstance(result, WorkRequest)
        self.assertEqual(result, work_request)
        self.assertEqual(result.worker, signing_worker)

    def assert_schedule_for_worker_return_none(
        self, status: WorkRequest.Statuses
    ) -> None:
        """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(self.playground.schedule_for_worker(self.worker))

    def test_schedule_for_worker_already_pending(self) -> None:
        """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) -> None:
        """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: mock.MagicMock,  # noqa: ARG002, U100
    ) -> None:
        """Consider a Worker's concurrency when assigning a WorkRequest."""
        worker = Worker.objects.get_or_create_celery()
        worker.concurrency = 3
        worker.dynamic_metadata = {
            "system:architectures": [native_architecture()],
            "system:worker_type": WorkerType.CELERY,
            "server:servernoop:version": ServerNoop.TASK_VERSION,
            "provided_tags": [
                wtags.WORKER_TYPE_CELERY,
                wtags.TASK_PREFIX + "server:servernoop:version:1",
            ],
        }

        for _ in range(3):
            work_request = self.playground.create_server_task(
                task_name="servernoop", mark_pending=True
            )
            self.assertEqual(
                work_request, self.playground.schedule_for_worker(worker)
            )

        work_request = self.playground.create_server_task(
            task_name="servernoop", mark_pending=True
        )
        self.assertIsNone(self.playground.schedule_for_worker(worker))

    def test_get_provided_worker_tags(self) -> None:
        worker = Worker.objects.get_or_create_celery()
        system_tags = {wtags.WORKER_TYPE_PREFIX + worker.worker_type}
        for worker_tags, admin_tags, expected in (
            ([], [], set()),
            (None, None, set()),
            (["a", "b"], [], {"a", "b"}),
            ([], ["a", "b"], {"a", "b"}),
            (["a"], ["b"], {"a", "b"}),
        ):
            with self.subTest(
                worker_tags=repr(worker_tags), admin_tags=repr(admin_tags)
            ):
                worker.dynamic_metadata = {"provided_tags": worker_tags}
                worker.static_metadata = {"provided_tags": admin_tags}
                self.assertEqual(
                    get_provided_worker_tags(worker), expected | system_tags
                )

    def test_get_provided_worker_tags_worker_type(self) -> None:
        self.worker.set_dynamic_metadata(
            {
                **self.worker.dynamic_metadata,
                "provided_tags": [wtags.WORKER_TYPE_SIGNING],
            }
        )
        # Worker type uses the system version, not the worker-provided one
        self.assertEqual(
            get_provided_worker_tags(self.worker), {wtags.WORKER_TYPE_EXTERNAL}
        )


class TagBasedScheduleForWorkerTests(ScheduleForWorkerTests):
    """Version of ScheduleForWorkerTests with tag based scheduling on."""

    @override
    def setUp(self) -> None:
        super().setUp()
        self.enterContext(override_settings(FEATURE_TAG_BASED_SCHEDULING=True))

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_task_required_tags_enforced(self) -> None:  # pragma: no cover
        work_request = self.create_sample_sbuild_work_request()
        for required, provided, accepted in (
            ([], [], True),
            ([], ["a"], True),
            (["a"], [], False),
            (["a"], ["a"], True),
            (["a"], ["b"], False),
            (["a", "b", "c"], ["a", "b"], False),
            (["a", "b", "c"], ["a", "b", "c"], True),
            (["a", "b", "c"], ["a", "b", "c", "d"], True),
        ):
            with (
                self.subTest(required=repr(required), provided=repr(provided)),
                mock.patch(
                    "debusine.server.scheduler.get_provided_worker_tags",
                    return_value=set(provided),
                ),
                mock.patch(
                    "debusine.server.scheduler"
                    "._assign_work_request_to_worker_or_error"
                ) as assign_work_request,
            ):
                work_request.scheduler_tags_required = required
                work_request.save()
                self.playground.schedule_for_worker(self.worker)
                if accepted:
                    assign_work_request.assert_called_with(
                        work_request, self.worker, set()
                    )
                else:
                    assign_work_request.assert_not_called()

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_worker_required_tags_enforced(self) -> None:  # pragma: no cover
        # Ensure the check of worker provided tags always passes here
        work_request = self.create_sample_sbuild_work_request(
            require_scheduler_tags=False
        )
        for required, provided, accepted in (
            ([], [], True),
            ([], ["a"], True),
            (["a"], [], False),
            (["a"], ["a"], True),
            (["a"], ["b"], False),
            (["a", "b", "c"], ["a", "b"], False),
            (["a", "b", "c"], ["a", "b", "c"], True),
            (["a", "b", "c"], ["a", "b", "c", "d"], True),
        ):
            with (
                self.subTest(required=repr(required), provided=repr(provided)),
                mock.patch(
                    "debusine.server.scheduler.get_required_task_tags",
                    return_value=set(required),
                ),
                mock.patch(
                    "debusine.server.scheduler"
                    "._assign_work_request_to_worker_or_error"
                ) as assign_work_request,
            ):
                work_request.scheduler_tags_provided = provided
                work_request.save()
                assert work_request.worker is None

                self.playground.schedule_for_worker(self.worker)

                if accepted:
                    assign_work_request.assert_called_with(
                        work_request, self.worker, set()
                    )
                else:
                    assign_work_request.assert_not_called()


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

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

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

        # Marking work_request_1 as pending scheduled a scheduler run, which
        # assigned it to a worker.
        await work_request_1.arefresh_from_db()
        work_request_1_worker = await database_sync_to_async(
            lambda: work_request_1.worker
        )()

        self.assertEqual(work_request_1_worker, self.worker)

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

        await work_request_2.arefresh_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
        await database_sync_to_async(work_request_1.mark_running)()
        await database_sync_to_async(work_request_1.mark_completed)(
            WorkRequest.Results.SUCCESS
        )

        await work_request_2.arefresh_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, CELERY_TASK_ALWAYS_EAGER=True
    )
    async def test_schedule_from_worker_changed(self) -> None:
        """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
        )(require_scheduler_tags=False)

        # 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 work_request_1.arefresh_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) -> None:
        """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(self.playground.schedule_for_worker(self.worker))
        finally:
            thread.stop_transaction()

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_when_work_request_is_locked(self) -> None:
        """schedule_for_worker() skips and logs work requests it can't lock."""
        workflow = self.playground.create_workflow()
        work_requests = [
            self.create_sample_sbuild_work_request(
                require_scheduler_tags=False, parent=parent
            )
            for parent in (workflow, workflow, None, None)
        ]

        thread = RunInParallelTransaction(
            lambda: WorkRequest.objects.filter(
                id__in={work_requests[0].id, work_requests[2].id}
            ).lock_workflows_for_update()
        )
        thread.start_transaction()

        try:
            with self.assertLogs(
                logger="debusine.server.scheduler", level=logging.DEBUG
            ) as logs:
                self.assertEqual(
                    self.playground.schedule_for_worker(self.worker),
                    work_requests[3],
                )
            # We skip work_requests[1], since we already found that another
            # work request in the same workflow was locked.
            self.assertEqual(
                logs.output,
                [
                    f"DEBUG:debusine.server.scheduler:"
                    f"[SCHEDULER] Failed to lock WorkRequest {work_request}"
                    f" (skipping this workflow)"
                    for work_request in (work_requests[0], work_requests[2])
                ],
            )
        finally:
            thread.stop_transaction()

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_when_work_request_is_locked_on_error(self) -> None:
        work_request = self.create_sample_sbuild_work_request(
            require_scheduler_tags=False
        )
        # invalid: it's missing required properties
        work_request.configured_task_data = {}
        work_request.save()

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

        try:
            with self.assertLogsContains(
                f"[SCHEDULER] Failed to lock WorkRequest {work_request}"
                f" to record error",
                logger="debusine.server.scheduler",
                level=logging.DEBUG,
            ):
                self.assertIsNone(
                    self.playground.schedule_for_worker(self.worker)
                )
        finally:
            thread.stop_transaction()

    def assert_dispatches_celery_task(
        self,
        mock_apply_async: mock.MagicMock,
        work_request: WorkRequest,
        task_id_prefix: str,
    ) -> None:
        """Assert that a work request is dispatched to a Celery task."""
        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"],
            rf"^{task_id_prefix}_{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 find_workflow_task_results(self) -> QuerySet[TaskResult]:
        return celery_tasks_in_progress(
            "debusine.server.workflows.celery.run_workflow_task"
        ).order_by("date_created")

    @override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
    def test_server_task(self) -> None:
        """schedule_for_worker() dispatches server tasks to Celery."""
        worker = Worker.objects.get_or_create_celery()
        worker.dynamic_metadata = {
            "system:architectures": [native_architecture()],
            "system:worker_type": WorkerType.CELERY,
            "server:servernoop:version": ServerNoop.TASK_VERSION,
            "provided_tags": [
                wtags.WORKER_TYPE_CELERY,
                wtags.TASK_PREFIX
                + f"server:servernoop:version:{ServerNoop.TASK_VERSION}",
            ],
        }
        work_request = self.playground.create_server_task(
            task_name="servernoop", mark_pending=True
        )

        with mock.patch(
            "debusine.server.scheduler.run_server_task.apply_async"
        ) as mock_apply_async:
            result = self.playground.schedule_for_worker(worker)

        self.assertEqual(work_request, result)
        self.assert_dispatches_celery_task(
            mock_apply_async, work_request, "servernoop"
        )

    def test_workflow_callbacks(self) -> None:
        """schedule() dispatches workflow callbacks to Celery."""
        WorkRequest.objects.all().delete()
        parent = self.playground.create_workflow()
        parent.mark_running()
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        wr.mark_pending()

        with mock.patch(
            "debusine.server.scheduler.run_workflow_task.apply_async"
        ) as mock_apply_async:
            self.assertEqual(schedule(), [wr])

        self.assert_dispatches_celery_task(
            mock_apply_async, wr, "internal_workflow"
        )
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.RUNNING)

    def test_workflow_callback_another_workflow_running(self) -> None:
        """Callbacks for multiple workflows may run concurrently."""
        test_app = TestApp(
            set_as_current=True,
            broker="memory://",
            task_cls=TaskWithBetterArgumentRepresentation,
        )
        self.enterContext(setup_default_app(test_app))

        WorkRequest.objects.all().delete()
        workflows = [self.playground.create_workflow() for _ in range(3)]
        for workflow in workflows:
            workflow.mark_running()
        wrs = [
            WorkRequest.objects.create_workflow_callback(
                parent=workflow, step="test"
            )
            for workflow in workflows[:2]
        ]
        for wr in wrs:
            wr.mark_pending()

        self.assertEqual(schedule(), wrs)

        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[wr.id] for wr in wrs[:2]],
        )
        self.assertQuerySetEqual(
            WorkRequest.objects.filter(id__in={wr.id for wr in wrs})
            .order_by("id")
            .values_list("status", flat=True),
            [WorkRequest.Statuses.RUNNING] * len(wrs),
        )

        wrs.append(
            WorkRequest.objects.create_workflow_callback(
                parent=workflows[-1], step="test"
            )
        )
        wrs[-1].mark_pending()
        # A running work request in the same parent workflow that isn't a
        # workflow callback doesn't stop the workflow callback running.
        self.playground.create_worker_task(
            parent=workflows[-1], assign_new_worker=True, mark_running=True
        )

        self.assertCountEqual(schedule(), [wrs[-1]])
        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[wr.id] for wr in wrs],
        )
        self.assertQuerySetEqual(
            WorkRequest.objects.filter(id__in={wr.id for wr in wrs})
            .order_by("id")
            .values_list("status", flat=True),
            [WorkRequest.Statuses.RUNNING] * len(wrs),
        )

    def test_workflow_callback_same_workflow_running(self) -> None:
        """Callbacks for the same workflow may not run concurrently."""
        test_app = TestApp(
            set_as_current=True,
            broker="memory://",
            task_cls=TaskWithBetterArgumentRepresentation,
        )
        self.enterContext(setup_default_app(test_app))

        WorkRequest.objects.all().delete()
        workflow = self.playground.create_workflow()
        workflow.mark_running()
        wrs = [
            WorkRequest.objects.create_workflow_callback(
                parent=workflow, step=f"test{i}"
            )
            for i in range(2)
        ]
        for wr in wrs:
            wr.mark_pending()

        self.assertEqual(schedule(), [wrs[0]])

        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[wrs[0].id]],
        )
        self.assertQuerySetEqual(
            WorkRequest.objects.filter(id__in={wr.id for wr in wrs})
            .order_by("id")
            .values_list("status", flat=True),
            [WorkRequest.Statuses.RUNNING, WorkRequest.Statuses.PENDING],
        )

        wrs.append(
            WorkRequest.objects.create_workflow_callback(
                parent=workflow, step="test2"
            )
        )
        wrs[-1].mark_pending()

        self.assertEqual(schedule(), [])

        workflow_task_results = self.find_workflow_task_results().filter(
            status=states.PENDING
        )
        self.assertQuerySetEqual(
            workflow_task_results.values_list("task_args_array", flat=True),
            [[wrs[0].id]],
        )
        self.assertQuerySetEqual(
            WorkRequest.objects.filter(id__in={wr.id for wr in wrs})
            .order_by("id")
            .values_list("status", flat=True),
            [
                WorkRequest.Statuses.RUNNING,
                WorkRequest.Statuses.PENDING,
                WorkRequest.Statuses.PENDING,
            ],
        )

        # Completing the Celery task (albeit artificially) allows the
        # scheduler to proceed.
        WorkRequest.objects.get(id=wrs[0].id).mark_completed(
            WorkRequest.Results.SUCCESS
        )
        workflow_task_result = workflow_task_results.get()
        workflow_task_result.status = states.SUCCESS
        workflow_task_result.save()

        self.assertEqual(schedule(), [wrs[1]])

        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[wrs[1].id]],
        )
        self.assertQuerySetEqual(
            WorkRequest.objects.filter(id__in={wr.id for wr in wrs})
            .order_by("id")
            .values_list("status", flat=True),
            [
                WorkRequest.Statuses.COMPLETED,
                WorkRequest.Statuses.RUNNING,
                WorkRequest.Statuses.PENDING,
            ],
        )

    def test_workflows(self) -> None:
        """schedule() dispatches workflows to Celery."""
        WorkRequest.objects.all().delete()
        wr = self.playground.create_workflow(task_name="noop")

        with mock.patch(
            "debusine.server.scheduler.run_workflow_task.apply_async"
        ) as mock_apply_async:
            self.assertEqual(schedule(), [wr])

        self.assert_dispatches_celery_task(
            mock_apply_async, wr, "workflow_noop"
        )
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.RUNNING)

    def test_sub_workflow(self) -> None:
        """schedule() orchestrates pending sub-workflows via their root."""
        root_wr = self.playground.create_workflow(task_name="noop")
        sub_wr = self.playground.create_workflow(parent=root_wr)
        # Mark the root workflow running, to simulate the situation where it
        # can only do part of its work and has to be called back later.
        root_wr.mark_running()

        with mock.patch(
            "debusine.server.scheduler.run_workflow_task.apply_async",
            side_effect=lambda *args, **kwargs: sub_wr.mark_running(),
        ) as mock_apply_async:
            schedule()

        self.assert_dispatches_celery_task(
            mock_apply_async, root_wr, "workflow_noop"
        )
        # The (fake) orchestrator was called and marked the sub-workflow as
        # running.
        self.assertEqual(sub_wr.status, WorkRequest.Statuses.RUNNING)

    def test_groups_workflows_by_root(self) -> None:
        """schedule() groups pending workflows by their root workflow."""
        root_wr = self.playground.create_workflow(task_name="noop")
        sub_wr = self.playground.create_workflow(parent=root_wr)

        with mock.patch(
            "debusine.server.scheduler.run_workflow_task.apply_async"
        ) as mock_apply_async:
            schedule()

        # Only the root workflow's orchestrator is called.
        self.assert_dispatches_celery_task(
            mock_apply_async, root_wr, "workflow_noop"
        )
        root_wr.refresh_from_db()
        sub_wr.refresh_from_db()
        self.assertEqual(root_wr.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(sub_wr.status, WorkRequest.Statuses.PENDING)

    def test_sub_workflow_another_workflow_running(self) -> None:
        """Sub-workflows of different root workflows may run concurrently."""
        test_app = TestApp(
            set_as_current=True,
            broker="memory://",
            task_cls=TaskWithBetterArgumentRepresentation,
        )
        self.enterContext(setup_default_app(test_app))

        WorkRequest.objects.all().delete()
        roots = [self.playground.create_workflow() for _ in range(3)]
        # Mark the root workflows running, to simulate the situation where
        # they can only do part of their work and have to be called back
        # later.
        for root in roots:
            root.mark_running()
        for root in roots[:2]:
            child = self.playground.create_workflow(parent=root)
            child.mark_running()
            self.playground.create_workflow(parent=child)

        self.assertEqual(schedule(), roots[:2])

        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[root.id] for root in roots[:2]],
        )

        self.playground.create_workflow(parent=roots[-1])
        # A running work request in the same parent workflow that isn't a
        # sub-workflow doesn't stop the sub-workflow running.
        self.playground.create_worker_task(
            parent=roots[-1], assign_new_worker=True, mark_running=True
        )

        self.assertCountEqual(schedule(), [roots[-1]])

        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[root.id] for root in roots],
        )

    def test_sub_workflow_same_workflow_running(self) -> None:
        """Sub-workflows of the same root workflow may not run concurrently."""
        test_app = TestApp(
            set_as_current=True,
            broker="memory://",
            task_cls=TaskWithBetterArgumentRepresentation,
        )
        self.enterContext(setup_default_app(test_app))

        WorkRequest.objects.all().delete()
        root = self.playground.create_workflow()
        root.mark_running()
        for _ in range(2):
            self.playground.create_workflow(parent=root)

        self.assertEqual(schedule(), [root])

        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[root.id]],
        )

        self.playground.create_workflow(parent=root)

        self.assertEqual(schedule(), [])

        workflow_task_results = self.find_workflow_task_results().filter(
            status=states.PENDING
        )
        self.assertQuerySetEqual(
            workflow_task_results.values_list("task_args_array", flat=True),
            [[root.id]],
        )

        # Completing the Celery task (albeit artificially) allows the
        # scheduler to proceed.
        workflow_task_result = workflow_task_results.get()
        workflow_task_result.status = states.SUCCESS
        workflow_task_result.save()

        self.assertEqual(schedule(), [root])

        self.assertQuerySetEqual(
            self.find_workflow_task_results()
            .filter(status=states.PENDING)
            .values_list("task_args_array", flat=True),
            [[root.id]],
        )


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

    @override_settings(PGTRANSACTION_RETRY=0)
    def test_updates_workflows(self) -> None:
        with transaction.atomic():
            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():
            schedule()

            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 function 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():
                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 = RunInParallelTransaction(create_workflow_needs_update)
            thread.run_and_wait()
            return mock.DEFAULT

        with mock.patch(
            "debusine.server.scheduler.workflow_ancestors",
            side_effect=create_workflow_needs_update_in_thread,
            wraps=workflow_ancestors,
        ):
            schedule()

        # 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 function'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)

    def test_concurrent_updates(self) -> None:
        """
        If the workflow lock query hits a serialization failure, the
        function aborts with an ``OperationalError``.
        """
        with transaction.atomic():
            workflow = self.playground.create_workflow()
            work_request = workflow.create_child_worker("noop")
            work_request.mark_aborted()

        with (
            self.assertLogsContains(
                "update_workflows() deferred due to a concurrent update",
                logger="debusine.server.scheduler",
                level=logging.WARNING,
            ),
            mock.patch(
                "debusine.db.models.work_requests.WorkRequestQuerySet.lock_workflows_for_update",
                side_effect=OperationalError(
                    "could not serialize access due to concurrent update"
                ),
            ),
        ):
            schedule()


class ScheduleSoonTests(TestCase):
    """Test schedule_soon()."""

    @override_settings(CELERY_TASK_ALWAYS_EAGER=True, PGTRANSACTION_RETRY=0)
    def test_runs(self) -> None:
        with (
            mock.patch("debusine.server.scheduler.schedule") as mock_schedule,
            self.captureOnCommitCallbacks(execute=True),
        ):
            schedule_soon()

        mock_schedule.assert_called_once()

    def test_batches(self) -> None:
        """schedule_soon() runs the scheduler once per savepoint."""
        with (
            transaction.atomic(),
            mock.patch(
                "debusine.server.scheduler.schedule_task.apply_async"
            ) as mock_apply_async,
            self.captureOnCommitCallbacks(execute=True) as callbacks,
        ):
            schedule_soon()
            schedule_soon()

        self.assertEqual(
            len(
                [
                    task
                    for task in get_apply_async_tasks(callbacks)
                    if task == schedule_task
                ]
            ),
            1,
        )
        mock_apply_async.assert_called_once()
