# Copyright 2019, 2021-2024 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Unit tests for the work request models."""

import io
import logging
from datetime import datetime, timedelta
from typing import Any, ClassVar
from unittest import mock

from asgiref.sync import sync_to_async

from django.core.exceptions import ValidationError
from django.db import transaction
from django.test import TestCase
from django.utils import timezone

try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic as pydantic  # type: ignore

from debusine.artifacts.models import ArtifactCategory, CollectionCategory
from debusine.db.models import (
    Collection,
    CollectionItem,
    Token,
    User,
    WorkRequest,
    Worker,
    WorkflowTemplate,
    Workspace,
    default_workspace,
)
from debusine.db.models.work_requests import CannotRetry
from debusine.db.models.workers import WorkerManager
from debusine.server.collections import (
    CollectionManagerInterface,
)
from debusine.server.management.commands.delete_expired import (
    DeleteExpiredWorkRequests,
    DeleteOperation,
)
from debusine.server.workflows import Workflow
from debusine.server.workflows.models import (
    BaseWorkflowData,
    WorkRequestWorkflowData,
)
from debusine.tasks.models import (
    ActionUpdateCollectionWithArtifacts,
    BaseDynamicTaskData,
    EventReactions,
    TaskTypes,
    WorkerType,
)
from debusine.tasks.noop import Noop
from debusine.tasks.server import TaskDatabaseInterface
from debusine.test import TestHelpersMixin
from debusine.test.django import ChannelsHelpersMixin
from debusine.test.utils import preserve_task_registry


class WorkerManagerTests(TestHelpersMixin, TestCase):
    """Tests for the WorkerManager."""

    def test_connected(self):
        """WorkerManager.connected() return the connected Workers."""
        worker_connected = Worker.objects.create_with_fqdn(
            'connected-worker', Token.objects.create()
        )

        worker_connected.mark_connected()

        Worker.objects.create_with_fqdn(
            'not-connected-worker',
            Token.objects.create(),
        )

        self.assertQuerysetEqual(Worker.objects.connected(), [worker_connected])

    def test_waiting_for_work_request(self):
        """Test WorkerManager.waiting_for_work_request() return a Worker."""
        worker = Worker.objects.create_with_fqdn(
            'worker-a', Token.objects.create(enabled=True)
        )

        # WorkerManagement.waiting_for_work_request: returns no workers because
        # the worker is not connected
        self.assertQuerysetEqual(Worker.objects.waiting_for_work_request(), [])

        worker.mark_connected()

        # Now the Worker is ready to have a task assigned
        self.assertQuerysetEqual(
            Worker.objects.waiting_for_work_request(), [worker]
        )

        # A task is assigned to the worker
        work_request = self.create_work_request(worker=worker)

        # The worker is not ready (it is busy with a task assigned)
        self.assertQuerysetEqual(Worker.objects.waiting_for_work_request(), [])

        # The task finished
        work_request.status = WorkRequest.Statuses.COMPLETED
        work_request.save()

        # The worker is ready: the WorkRequest that had assigned finished
        self.assertQuerysetEqual(
            Worker.objects.waiting_for_work_request(), [worker]
        )

    def test_waiting_for_work_request_celery(self):
        """
        Test WorkerManager.waiting_for_work_request() with Celery workers.

        Celery workers have no tokens, so they require special handling.
        """
        worker = Worker.objects.get_or_create_celery()
        worker.concurrency = 3
        worker.save()

        # WorkerManagement.waiting_for_work_request: returns no workers because
        # the worker is not connected
        self.assertQuerysetEqual(Worker.objects.waiting_for_work_request(), [])

        worker.mark_connected()

        # Now the Worker is ready to have a task assigned
        self.assertQuerysetEqual(
            Worker.objects.waiting_for_work_request(), [worker]
        )

        # Assign two tasks to the worker; it can still accept more work
        for _ in range(2):
            work_request = self.create_work_request(worker=worker)
            self.assertQuerysetEqual(
                Worker.objects.waiting_for_work_request(), [worker]
            )

        # Assign a third task to the worker; it is now as busy as it allows
        work_request = self.create_work_request(worker=worker)
        self.assertQuerysetEqual(Worker.objects.waiting_for_work_request(), [])

        # The third task finished
        work_request.status = WorkRequest.Statuses.COMPLETED
        work_request.save()

        # Although it is still running some tasks, the worker can accept
        # more work again
        self.assertQuerysetEqual(
            Worker.objects.waiting_for_work_request(), [worker]
        )

    def test_waiting_for_work_request_no_return_disabled_workers(self):
        """Test WorkerManager.waiting_for_work_request() no return disabled."""
        worker_enabled = Worker.objects.create_with_fqdn(
            "worker-enabled", Token.objects.create(enabled=True)
        )
        worker_enabled.mark_connected()

        worker_disabled = Worker.objects.create_with_fqdn(
            "worker-disabled", Token.objects.create(enabled=False)
        )
        worker_disabled.mark_connected()

        self.assertQuerysetEqual(
            Worker.objects.waiting_for_work_request(), [worker_enabled]
        )

    def test_create_with_fqdn_new_fqdn(self):
        """WorkerManager.create_with_fqdn() return a worker."""
        token = Token.objects.create()
        worker = Worker.objects.create_with_fqdn(
            'a-new-and-unique-name', token=token
        )

        self.assertEqual(worker.name, 'a-new-and-unique-name')
        self.assertEqual(worker.token, token)
        self.assertIsNotNone(worker.pk)

    def test_create_with_fqdn_duplicate_fqdn(self):
        """
        WorkerManager.create_with_fqdn() return a worker.

        The name ends with -2 because 'connected-worker' is already used.
        """
        Worker.objects.create_with_fqdn(
            'connected-worker', token=Token.objects.create()
        )

        token = Token.objects.create()
        worker = Worker.objects.create_with_fqdn(
            'connected-worker', token=token
        )

        self.assertEqual(worker.name, 'connected-worker-2')
        self.assertEqual(worker.token, token)
        self.assertEqual(worker.worker_type, WorkerType.EXTERNAL)
        self.assertIsNotNone(worker.pk)

    def test_get_or_create_celery(self):
        """WorkerManager.get_or_create_celery returns a Celery worker."""
        worker = Worker.objects.get_or_create_celery()

        self.assertEqual(worker.name, "celery")
        self.assertIsNone(worker.token)
        self.assertEqual(worker.worker_type, WorkerType.CELERY)
        self.assertIsNotNone(worker.pk)

        self.assertEqual(Worker.objects.get_or_create_celery(), worker)

    def test_slugify_with_suffix_counter_1(self):
        """WorkerManager._generate_unique_name does not append '-1'."""
        self.assertEqual(
            WorkerManager._generate_unique_name('worker.lan', 1), 'worker-lan'
        )

    def test_slugify_with_suffix_counter_3(self):
        """WorkerManager._generate_unique_name appends '-3'."""
        self.assertEqual(
            WorkerManager._generate_unique_name('worker.lan', 3), 'worker-lan-3'
        )

    def test_get_worker_by_token_or_none_return_none(self):
        """WorkerManager.get_worker_by_token_or_none() return None."""
        self.assertIsNone(
            Worker.objects.get_worker_by_token_key_or_none('non-existing-key')
        )

    def test_get_worker_by_token_or_none_return_worker(self):
        """WorkerManager.get_worker_by_token_or_none() return the Worker."""
        token = Token.objects.create()

        worker = Worker.objects.create_with_fqdn('worker-a', token)

        self.assertEqual(
            Worker.objects.get_worker_by_token_key_or_none(token.key), worker
        )

    def test_get_worker_or_none_return_worker(self):
        """WorkerManager.get_worker_or_none() return the Worker."""
        token = Token.objects.create()

        worker = Worker.objects.create_with_fqdn('worker-a', token)

        self.assertEqual(Worker.objects.get_worker_or_none('worker-a'), worker)

    def test_get_worker_or_none_return_none(self):
        """WorkerManager.get_worker_or_none() return None."""
        self.assertIsNone(Worker.objects.get_worker_or_none('does-not-exist'))


class WorkerTests(TestHelpersMixin, TestCase):
    """Tests for the Worker model."""

    worker: ClassVar[Worker]

    @classmethod
    def setUpTestData(cls):
        """Set up the Worker for the tests."""
        super().setUpTestData()
        cls.worker = Worker.objects.create_with_fqdn(
            "computer.lan", Token.objects.create()
        )
        cls.worker.static_metadata = {"os": "debian"}
        cls.worker.set_dynamic_metadata({"cpu_cores": "4"})
        cls.worker.save()

    def test_mark_connected(self):
        """Test mark_connect method."""
        time_before = timezone.now()
        self.assertIsNone(self.worker.connected_at)

        self.worker.mark_connected()

        self.assertGreaterEqual(self.worker.connected_at, time_before)
        self.assertLessEqual(self.worker.connected_at, timezone.now())

    def test_mark_disconnected(self):
        """Test mark_disconnected method."""
        self.worker.mark_connected()

        self.assertTrue(self.worker.connected())
        self.worker.mark_disconnected()

        self.assertFalse(self.worker.connected())
        self.assertIsNone(self.worker.connected_at)

    def test_connected(self):
        """Test connected method."""
        self.assertFalse(self.worker.connected())

        self.worker.connected_at = timezone.now()

        self.assertTrue(self.worker.connected())

    def test_is_busy(self):
        """Test is_busy method."""
        worker = Worker.objects.create_with_fqdn(
            "test", token=Token.objects.create()
        )
        self.assertFalse(worker.is_busy())

        work_request = self.create_work_request(worker=worker)

        self.assertEqual(work_request.status, WorkRequest.Statuses.PENDING)
        self.assertTrue(worker.is_busy())

        work_request.assign_worker(worker)

        self.assertTrue(worker.is_busy())

        work_request.mark_running()

        self.assertTrue(worker.is_busy())

        work_request.mark_aborted()

        self.assertFalse(worker.is_busy())

    def test_is_busy_concurrency(self):
        """The is_busy method handles workers with concurrency > 1."""
        worker = Worker.objects.get_or_create_celery()
        worker.concurrency = 3
        worker.save()
        for _ in range(3):
            self.assertFalse(worker.is_busy())
            work_request = self.create_work_request(worker=worker)
            work_request.assign_worker(worker)
        self.assertTrue(worker.is_busy())
        work_request.mark_aborted()
        self.assertFalse(worker.is_busy())

    def test_metadata_no_conflict(self):
        """Test metadata method: return all the metadata."""
        self.assertEqual(
            self.worker.metadata(), {'cpu_cores': '4', 'os': 'debian'}
        )

    def test_metadata_with_conflict(self):
        """
        Test metadata method: return all the metadata.

        static_metadata has priority over dynamic_metadata
        """
        # Assert initial state
        self.assertEqual(self.worker.dynamic_metadata['cpu_cores'], '4')

        # Add new static_metadata key
        self.worker.static_metadata['cpu_cores'] = '8'

        self.assertEqual(
            self.worker.metadata(), {'cpu_cores': '8', 'os': 'debian'}
        )

    def test_metadata_is_deep_copy(self):
        """Test metadata does a deep copy."""
        self.worker.dynamic_metadata['schroots'] = ['buster', 'bullseye']
        self.worker.metadata()['schroots'].append('bookworm')

        self.assertEqual(
            self.worker.dynamic_metadata['schroots'], ['buster', 'bullseye']
        )

    def test_set_dynamic_metadata(self):
        """Worker.set_dynamic_metadata sets the dynamic metadata."""
        self.worker.dynamic_metadata = {}
        self.worker.dynamic_metadata_updated_at = None
        self.worker.save()

        dynamic_metadata = {"cpu_cores": 4, "ram": "16"}
        self.worker.set_dynamic_metadata(dynamic_metadata)

        self.worker.refresh_from_db()

        self.assertEqual(self.worker.dynamic_metadata, dynamic_metadata)
        self.assertLessEqual(
            self.worker.dynamic_metadata_updated_at, timezone.now()
        )

    def test_str(self):
        """Test WorkerTests.__str__."""
        self.assertEqual(
            self.worker.__str__(),
            f"Id: {self.worker.id} Name: {self.worker.name}",
        )


class TestWorkflowData(BaseWorkflowData):
    """Workflow data to test instantiation."""

    ivar: int = 1
    dvar: dict[str, Any] = pydantic.Field(default_factory=dict)


class WorkRequestManagerTestsTestWorkflow(Workflow[TestWorkflowData]):
    """Workflow used to test instantiation of data."""

    # This needs a unique name to avoid conflicting with other tests, because
    # it gets registered as a Workflow subclass

    def populate(self) -> None:
        """Unused abstract method from Workflow."""
        raise NotImplementedError()


class WorkRequestManagerTests(TestHelpersMixin, TestCase):
    """Tests for WorkRequestManager."""

    def test_create_workflow(self) -> None:
        """Workflows are created correctly."""
        user = self.get_test_user()

        workflow_name = "workrequestmanagerteststestworkflow"

        template = WorkflowTemplate.objects.create(
            name="test",
            workspace=default_workspace(),
            task_name=workflow_name,
            task_data={},
            priority=10,
        )

        wr = WorkRequest.objects.create_workflow(
            template=template, created_by=user, data={}
        )
        self.assertEqual(wr.workspace, template.workspace)
        self.assertIsNone(wr.parent)
        self.assertEqual(wr.created_by, user)
        self.assertEqual(wr.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(wr.task_type, TaskTypes.WORKFLOW)
        self.assertEqual(wr.task_name, workflow_name)
        self.assertEqual(wr.task_data, {})
        self.assertEqual(wr.priority_base, 10)
        self.assertQuerysetEqual(wr.dependencies.all(), [])
        assert wr.internal_collection is not None
        self.assertEqual(wr.internal_collection.name, f"workflow-{wr.id}")
        self.assertEqual(
            wr.internal_collection.category,
            CollectionCategory.WORKFLOW_INTERNAL,
        )
        self.assertEqual(wr.internal_collection.workspace, template.workspace)
        self.assertEqual(
            wr.internal_collection.retains_artifacts,
            Collection.RetainsArtifacts.WORKFLOW,
        )
        self.assertEqual(wr.internal_collection.workflow, wr)

    def test_create_workflow_with_parent(self) -> None:
        """A sub-workflow does not have its own internal collection."""
        user = self.get_test_user()
        template = WorkflowTemplate.objects.create(
            name="test",
            workspace=default_workspace(),
            task_name="workrequestmanagerteststestworkflow",
        )
        root = WorkRequest.objects.create_workflow(
            template=template, created_by=user, data={}
        )
        sub_template = WorkflowTemplate.objects.create(
            name="subtest", workspace=default_workspace(), task_name="noop"
        )
        sub = WorkRequest.objects.create_workflow(
            template=sub_template, created_by=user, data={}, parent=root
        )

        self.assertEqual(sub.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(sub.parent, root)
        self.assertQuerysetEqual(sub.dependencies.all(), [])
        self.assertIsNone(sub.internal_collection)

    def test_create_workflow_with_parms(self) -> None:
        """Check that WorkflowTemplate parameters have precedence."""
        user = self.get_test_user()

        workflow_name = "workrequestmanagerteststestworkflow"

        template = WorkflowTemplate.objects.create(
            name="test",
            workspace=default_workspace(),
            task_name=workflow_name,
            task_data={"ivar": 2, "dvar": {"a": 1}},
        )

        wr = WorkRequest.objects.create_workflow(
            template=template, created_by=user, data={}
        )
        self.assertEqual(wr.task_data, {"ivar": 2, "dvar": {"a": 1}})

        wr = WorkRequest.objects.create_workflow(
            template=template,
            created_by=user,
            data={"ivar": 3, "dvar": {"b": 2}},
        )
        self.assertEqual(wr.task_data, {"ivar": 2, "dvar": {"a": 1}})

        template.task_data = {"ivar": 1}
        wr = WorkRequest.objects.create_workflow(
            template=template,
            created_by=user,
            data={"ivar": 3, "dvar": {"b": 2}},
        )
        self.assertEqual(wr.task_data, {"ivar": 1, "dvar": {"b": 2}})

    def test_create_workflow_validation(self) -> None:
        """Check validation of parameters on workflow instantiation."""
        user = self.get_test_user()

        workflow_name = "workrequestmanagerteststestworkflow"

        template = WorkflowTemplate.objects.create(
            name="test",
            workspace=default_workspace(),
            task_name=workflow_name,
            task_data={},
        )

        with self.assertLogsContains(
            "Cannot create a workflow",
            logger="debusine.db.models",
            level=logging.ERROR,
        ):
            wr = WorkRequest.objects.create_workflow(
                template=template,
                created_by=user,
                data={"noise": 1},
            )
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_create_synchronization_point(self) -> None:
        """Synchronization points are created correctly."""
        user = self.get_test_user()

        parent = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=user,
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
            task_data={},
            priority_base=5,
            priority_adjustment=1,
        )

        wr = WorkRequest.objects.create_synchronization_point(
            parent=parent, step="test"
        )
        self.assertEqual(wr.workspace, parent.workspace)
        self.assertEqual(wr.parent, parent)
        self.assertIsNotNone(wr.created_at)
        self.assertIsNone(wr.started_at)
        self.assertIsNone(wr.completed_at)
        self.assertEqual(wr.created_by, user)
        self.assertEqual(wr.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(wr.result, WorkRequest.Results.NONE)
        self.assertIsNone(wr.worker)
        self.assertEqual(wr.task_type, TaskTypes.INTERNAL)
        self.assertEqual(wr.task_name, "synchronization_point")
        self.assertEqual(wr.task_data, {})
        self.assertEqual(wr.priority_base, 6)
        self.assertQuerysetEqual(wr.dependencies.all(), [])
        self.assertEqual(
            wr.workflow_data,
            WorkRequestWorkflowData(step="test", display_name="test"),
        )
        wr.full_clean()

        workspace = Workspace.objects.create_with_name(name="test")

        parent = WorkRequest.objects.create(
            workspace=workspace,
            created_by=user,
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
            task_data={},
        )
        wr = WorkRequest.objects.create_synchronization_point(
            parent=parent,
            step="test",
            display_name="Test",
            status=WorkRequest.Statuses.PENDING,
        )
        self.assertEqual(wr.workspace, workspace)
        self.assertEqual(wr.parent, parent)
        self.assertIsNotNone(wr.created_at)
        self.assertIsNone(wr.started_at)
        self.assertIsNone(wr.completed_at)
        self.assertEqual(wr.created_by, user)
        self.assertEqual(wr.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(wr.result, WorkRequest.Results.NONE)
        self.assertIsNone(wr.worker)
        self.assertEqual(wr.task_type, TaskTypes.INTERNAL)
        self.assertEqual(wr.task_name, "synchronization_point")
        self.assertEqual(wr.task_data, {})
        self.assertEqual(
            wr.workflow_data,
            WorkRequestWorkflowData(step="test", display_name="Test"),
        )
        wr.full_clean()

    def test_create_workflow_callback(self) -> None:
        """Workflow callbacks are created correctly."""
        user = self.get_test_user()

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

        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        self.assertEqual(wr.workspace, parent.workspace)
        self.assertEqual(wr.parent, parent)
        self.assertIsNotNone(wr.created_at)
        self.assertIsNone(wr.started_at)
        self.assertIsNone(wr.completed_at)
        self.assertEqual(wr.created_by, user)
        self.assertEqual(wr.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(wr.result, WorkRequest.Results.NONE)
        self.assertIsNone(wr.worker)
        self.assertEqual(wr.task_type, TaskTypes.INTERNAL)
        self.assertEqual(wr.task_name, "workflow")
        self.assertEqual(wr.task_data, {})
        self.assertEqual(wr.priority_base, 0)
        self.assertQuerysetEqual(wr.dependencies.all(), [])
        self.assertEqual(
            wr.workflow_data,
            WorkRequestWorkflowData(step="test", display_name="test"),
        )
        wr.full_clean()

        workspace = Workspace.objects.create_with_name(name="test")

        parent = WorkRequest.objects.create(
            workspace=workspace,
            created_by=user,
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
            task_data={},
            priority_base=20,
            priority_adjustment=10,
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent,
            status=WorkRequest.Statuses.PENDING,
            step="test1",
            display_name="Test One",
        )
        self.assertEqual(wr.workspace, workspace)
        self.assertEqual(wr.parent, parent)
        self.assertIsNotNone(wr.created_at)
        self.assertIsNone(wr.started_at)
        self.assertIsNone(wr.completed_at)
        self.assertEqual(wr.created_by, user)
        self.assertEqual(wr.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(wr.result, WorkRequest.Results.NONE)
        self.assertIsNone(wr.worker)
        self.assertEqual(wr.task_type, TaskTypes.INTERNAL)
        self.assertEqual(wr.task_name, "workflow")
        self.assertEqual(wr.task_data, {})
        self.assertEqual(wr.priority_base, 30)
        self.assertEqual(
            wr.workflow_data,
            WorkRequestWorkflowData(step="test1", display_name="Test One"),
        )
        wr.full_clean()

    def test_pending(self):
        """WorkRequestManager.pending() returns pending WorkRequests."""
        work_request_1 = self.create_work_request(
            status=WorkRequest.Statuses.PENDING
        )

        work_request_2 = self.create_work_request(
            status=WorkRequest.Statuses.PENDING
        )

        self.create_work_request(status=WorkRequest.Statuses.ABORTED)

        self.assertQuerysetEqual(
            WorkRequest.objects.pending(), [work_request_1, work_request_2]
        )

        work_request_1.created_at = timezone.now()
        work_request_1.save()

        self.assertQuerysetEqual(
            WorkRequest.objects.pending(), [work_request_2, work_request_1]
        )

        work_request_1.priority_adjustment = 1
        work_request_1.save()

        self.assertQuerysetEqual(
            WorkRequest.objects.pending(), [work_request_1, work_request_2]
        )

        work_request_2.priority_base = 2
        work_request_2.save()

        self.assertQuerysetEqual(
            WorkRequest.objects.pending(), [work_request_2, work_request_1]
        )

    def test_pending_filter_by_worker(self):
        """WorkRequestManager.pending() returns WorkRequest for the worker."""
        self.create_work_request(
            status=WorkRequest.Statuses.PENDING, task_name="sbuild"
        )

        worker = Worker.objects.create_with_fqdn(
            "computer.lan", Token.objects.create()
        )

        work_request_2 = self.create_work_request(
            status=WorkRequest.Statuses.PENDING, worker=worker
        )

        self.assertQuerysetEqual(
            WorkRequest.objects.pending(worker=worker), [work_request_2]
        )

    def test_raise_value_error_exclude_assigned_and_worker(self):
        """WorkRequestManager.pending() raises ValueError."""
        worker = Worker.objects.create_with_fqdn(
            "computer.lan", Token.objects.create()
        )

        with self.assertRaisesRegex(
            ValueError, "Cannot exclude_assigned and filter by worker"
        ):
            WorkRequest.objects.pending(exclude_assigned=True, worker=worker)

    def test_pending_exclude_assigned(self):
        """
        Test WorkRequestManager.pending(exclude_assigned=True).

        It excludes work requests that are assigned to a worker.
        """
        # Pending, not assigned to a worker WorkRequest
        work_request = self.create_work_request(
            status=WorkRequest.Statuses.PENDING, task_name="sbuild"
        )

        # Is returned as expected
        self.assertQuerysetEqual(
            WorkRequest.objects.pending(exclude_assigned=True), [work_request]
        )

        # Creates a worker
        worker = Worker.objects.create_with_fqdn(
            'test', token=Token.objects.create()
        )

        # Assigns the worker to the work_request
        work_request.worker = worker
        work_request.save()

        # pending(exclude_assigned=True) doesn't return it anymore
        self.assertQuerysetEqual(
            WorkRequest.objects.pending(exclude_assigned=True), []
        )

        # without the exclude_assigned it returns it
        self.assertQuerysetEqual(
            WorkRequest.objects.pending(exclude_assigned=False), [work_request]
        )

    def test_running(self):
        """WorkRequestManager.running() returns running WorkRequests."""
        work_request = self.create_work_request(
            status=WorkRequest.Statuses.RUNNING
        )
        self.create_work_request(status=WorkRequest.Statuses.ABORTED)

        self.assertQuerysetEqual(WorkRequest.objects.running(), [work_request])

    def test_completed(self):
        """WorkRequestManager.completed() returns completed WorkRequests."""
        work_request = self.create_work_request(
            status=WorkRequest.Statuses.COMPLETED
        )
        self.create_work_request(status=WorkRequest.Statuses.RUNNING)

        self.assertQuerysetEqual(
            WorkRequest.objects.completed(), [work_request]
        )

    def test_aborted(self):
        """WorkRequestManager.aborted() returns aborted WorkRequests."""
        work_request = self.create_work_request(
            status=WorkRequest.Statuses.ABORTED
        )
        self.create_work_request(status=WorkRequest.Statuses.RUNNING)

        self.assertQuerysetEqual(WorkRequest.objects.aborted(), [work_request])

    def test_aborted_or_failed(self):
        """WorkRequestManager.aborted() returns aborted WorkRequests."""
        wr_aborted = self.create_work_request(
            status=WorkRequest.Statuses.ABORTED
        )
        wr_failed = self.create_work_request(
            result=WorkRequest.Results.FAILURE,
        )
        self.create_work_request(status=WorkRequest.Statuses.RUNNING)

        self.assertQuerysetEqual(
            WorkRequest.objects.aborted_or_failed(), [wr_aborted, wr_failed]
        )

    def test_expired(self):
        """Test expired() method."""
        wr = self.create_work_request(task_name="noop")

        wr.created_at = timezone.now() - timedelta(days=2)
        wr.expiration_delay = timedelta(days=1)
        wr.save()
        self.assertQuerysetEqual(
            WorkRequest.objects.expired(timezone.now()), [wr]
        )

        wr.expiration_delay = timedelta(0)
        wr.save()
        self.assertQuerysetEqual(
            WorkRequest.objects.expired(timezone.now()), []
        )

        wr.created_at = timezone.now()
        wr.expiration_delay = timedelta(days=1)
        wr.save()
        self.assertQuerysetEqual(
            WorkRequest.objects.expired(timezone.now()), []
        )

        workspace = Workspace.objects.create_with_name(name="test")
        workspace.default_expiration_delay = timedelta(days=2)
        workspace.save()
        wr.workspace = workspace
        wr.created_at = timezone.now() - timedelta(days=3)
        wr.expiration_delay = None
        wr.save()
        self.assertQuerysetEqual(
            WorkRequest.objects.expired(timezone.now()), [wr]
        )

        workspace.default_expiration_delay = timedelta(days=5)
        workspace.save()
        self.assertQuerysetEqual(
            WorkRequest.objects.expired(timezone.now()), []
        )

        workspace.default_expiration_delay = timedelta(days=0)
        workspace.save()
        self.assertQuerysetEqual(
            WorkRequest.objects.expired(timezone.now()), []
        )


class WorkRequestTests(TestHelpersMixin, TestCase):
    """Tests for the WorkRequest class."""

    user: ClassVar[User]
    template: ClassVar[WorkflowTemplate]
    workflow: ClassVar[WorkRequest]

    @classmethod
    def setUpTestData(cls) -> None:
        """Add a sample user to the test fixture."""
        super().setUpTestData()
        cls.user = cls.playground.get_default_user()
        cls.template = cls.playground.create_workflow_template(
            name="test",
            task_name="noop",
        )
        cls.workflow = WorkRequest.objects.create_workflow(
            template=cls.template, data={}, created_by=cls.user
        )

    def assert_data_must_be_dictionary(
        self, wr: WorkRequest, messages: list[str] | None = None
    ):
        """Check that a WorkRequest's task_data is a dictionary."""
        if messages is None:
            messages = ["task data must be a dictionary"]
        with self.assertRaises(ValidationError) as raised:
            wr.full_clean()
        self.assertEqual(
            raised.exception.message_dict,
            {"task_data": messages},
        )

    def test_clean_task_data_must_be_dict(self) -> None:
        """Ensure that task_data is a dict."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.user,
            task_name="noop",
            task_data={},
        )
        wr.full_clean()

        wr.task_data = None
        self.assert_data_must_be_dictionary(wr)

        wr.task_data = ""
        self.assert_data_must_be_dictionary(wr)

        wr.task_data = 3
        self.assert_data_must_be_dictionary(wr)

        wr.task_data = []
        self.assert_data_must_be_dictionary(wr)

        wr.task_data = object()
        self.assert_data_must_be_dictionary(
            wr,
            messages=[
                'Value must be valid JSON.',
                'task data must be a dictionary',
            ],
        )

    def test_clean_task_name_must_be_valid_worker(self) -> None:
        """Check validation of worker task names."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.user,
            task_name="noop",
            task_data={},
        )
        wr.full_clean()

        wr.task_name = "does-not-exist"
        with self.assertRaises(ValidationError) as raised:
            wr.full_clean()

        self.assertEqual(
            raised.exception.message_dict,
            {"task_name": ['does-not-exist: invalid Worker task name']},
        )

        wr.task_name = "does_not_exist"
        with self.assertRaises(ValidationError) as raised:
            wr.full_clean()

        self.assertEqual(
            raised.exception.message_dict,
            {
                "task_name": [
                    "does_not_exist: invalid Worker task name",
                ]
            },
        )

    def test_clean_task_name_must_be_valid_server(self) -> None:
        """Check validation of server task names."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.user,
            task_type=TaskTypes.SERVER,
            task_name="servernoop",
            task_data={},
        )
        wr.full_clean()

        wr.task_name = "does-not-exist"
        with self.assertRaises(ValidationError) as raised:
            wr.full_clean()

        self.assertEqual(
            raised.exception.message_dict,
            {"task_name": ['does-not-exist: invalid Server task name']},
        )

    def test_clean_task_name_must_be_valid_internal(self) -> None:
        """Check validation of internal task names."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.user,
            task_type=TaskTypes.INTERNAL,
            task_name="synchronization_point",
            task_data={},
        )
        wr.full_clean()

        wr.task_name = "workflow"
        wr.full_clean()

        wr.task_name = "does_not_exist"
        with self.assertRaises(ValidationError) as raised:
            wr.full_clean()
        self.assertEqual(
            raised.exception.message_dict,
            {
                "task_name": [
                    "does_not_exist: invalid task name for internal task"
                ]
            },
        )

    def test_clean_task_name_must_be_valid_workflow(self) -> None:
        """Check validation of workflow task names."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.user,
            task_type=TaskTypes.WORKFLOW,
            task_name="servernoop",
            task_data={},
        )
        with self.assertRaises(ValidationError) as raised:
            wr.full_clean()
        self.assertEqual(
            raised.exception.message_dict,
            {"task_name": ["servernoop: invalid workflow name"]},
        )

        wr.task_name = "noop"
        wr.full_clean()

    def test_clean_task_name_must_be_valid_signing(self) -> None:
        """Check validation of signing task names."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.user,
            task_type=TaskTypes.SIGNING,
            task_name="generatekey",
            task_data={"purpose": "uefi", "description": "A UEFI key"},
        )
        wr.full_clean()

        wr.task_name = "does-not-exist"
        with self.assertRaises(ValidationError) as raised:
            wr.full_clean()

        self.assertEqual(
            raised.exception.message_dict,
            {"task_name": ['does-not-exist: invalid Signing task name']},
        )

    def test_clean_invalid_task_type(self):
        """An invalid work request task type is rejected."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.user,
            task_type="invalid",
            task_name="noop",
            task_data={},
        )
        with self.assertRaises(NotImplementedError):
            wr.full_clean()

    def test_clean_data_task_data_raise_validation_error(self) -> None:
        """Check validation of task data."""

        def make_task_data():
            return {
                "build_components": ["any", "all"],
                "host_architecture": "amd64",
                "environment": (
                    "debian/match:codename=bookworm:architecture=amd64"
                ),
                "input": {
                    "source_artifact": 5,
                },
            }

        wr = WorkRequest.objects.create(
            created_by=self.user,
            workspace=default_workspace(),
            task_name="sbuild",
            task_data=make_task_data(),
        )
        wr.full_clean()

        wr.task_data = make_task_data()
        wr.task_data["invalid-key"] = None
        with self.assertRaisesRegex(
            ValidationError,
            "extra fields not permitted",
        ):
            wr.full_clean()

        wr.task_data = make_task_data()
        del wr.task_data["environment"]
        with self.assertRaisesRegex(
            ValidationError,
            'The backend "auto" requires "environment" to be set',
        ):
            wr.full_clean()

    def test_clean_workflow_data_raise_validation_error(self) -> None:
        """Check validation of workflow data."""
        wr = WorkRequest.objects.create(
            created_by=self.user,
            workspace=default_workspace(),
            task_type=TaskTypes.WORKFLOW,
            task_name="noop",
            task_data={},
        )
        wr.full_clean()

        wr.task_data["spurious"] = True
        with self.assertRaisesRegex(
            ValidationError,
            "extra fields not permitted",
        ):
            wr.full_clean()

    def test_workflow_data_accessors(self) -> None:
        """Test the custom pydantic accessors for workflow_data."""
        common_args = {
            "created_by": self.user,
            "workspace": default_workspace(),
            "task_type": TaskTypes.WORKER,
            "task_name": "noop",
            "task_data": {},
        }

        # Empty workflow_data
        wr = WorkRequest.objects.create(**common_args)
        self.assertEqual(wr.workflow_data_json, {})
        self.assertIsNone(wr.workflow_data)

        wr = WorkRequest.objects.create(**common_args)
        self.assertEqual(wr.workflow_data_json, {})
        self.assertIsNone(wr.workflow_data)

        # Filled workflow data
        wr = WorkRequest.objects.create(
            workflow_data_json={"display_name": "Test", "step": "test"},
            **common_args,
        )
        self.assertEqual(
            wr.workflow_data_json,
            {"display_name": "Test", "step": "test"},
        )
        self.assertEqual(
            wr.workflow_data,
            WorkRequestWorkflowData(display_name="Test", step="test"),
        )

        wr.workflow_data = None
        self.assertEqual(wr.workflow_data_json, {})
        self.assertIsNone(wr.workflow_data)

        wr.workflow_data = WorkRequestWorkflowData(
            display_name="Test1", step="test1"
        )
        self.assertEqual(
            wr.workflow_data_json,
            {"display_name": "Test1", "step": "test1"},
        )
        self.assertEqual(
            wr.workflow_data,
            WorkRequestWorkflowData(display_name="Test1", step="test1"),
        )

    def test_supersedes(self) -> None:
        """Test behaviour of the supersedes field."""
        wr = self.create_work_request(task_name="noop")
        self.assertIsNone(wr.supersedes)
        self.assertFalse(hasattr(wr, "superseded"))

        wr1 = self.create_work_request(task_name="noop")
        wr1.supersedes = wr
        wr1.save()

        wr.refresh_from_db()
        wr1.refresh_from_db()

        self.assertIsNone(wr.supersedes)
        self.assertEqual(wr1.supersedes, wr)
        self.assertEqual(wr.superseded, wr1)
        self.assertFalse(hasattr(wr1, "superseded"))

        self.assertQuerysetEqual(
            WorkRequest.objects.filter(superseded__isnull=False), [wr]
        )
        self.assertQuerysetEqual(
            WorkRequest.objects.filter(supersedes__isnull=False), [wr1]
        )

    def test_supersedes_expired(self) -> None:
        """Test expiring a superseded work request."""
        wr = self.create_work_request(task_name="noop")
        wr.created_at = timezone.now() - timedelta(days=365)
        wr.expiration_delay = timedelta(days=7)
        wr.save()
        wr1 = self.create_work_request(task_name="noop")
        wr1.supersedes = wr
        wr1.save()

        # Expire the superseded work request
        with (
            io.StringIO() as stdout,
            io.StringIO() as stderr,
            DeleteOperation(out=stdout, err=stderr, dry_run=False) as operation,
        ):
            delete_expired = DeleteExpiredWorkRequests(operation)
            delete_expired.run()

        self.assertQuerysetEqual(WorkRequest.objects.filter(pk=wr.pk), [])
        wr1.refresh_from_db()

        # Supersedes is set to None
        self.assertIsNone(wr1.supersedes)

    def test_cant_retry_superseded(self) -> None:
        """Superseded work requests cannot be retried."""
        wr = self.workflow.create_child("noop")
        wr.status = WorkRequest.Statuses.ABORTED
        wr.save()
        self.assertTrue(wr.can_retry())

        wr1 = self.workflow.create_child("noop")
        wr1.supersedes = wr
        self.assertFalse(wr.can_retry())
        with self.assertRaisesRegex(
            CannotRetry, r"^Cannot retry old superseded tasks$"
        ):
            wr.retry()

    def test_can_retry_aborted_failed(self) -> None:
        """Only aborted/failed work requests can be retried."""
        wr = self.workflow.create_child("noop")
        wr.status = WorkRequest.Statuses.COMPLETED
        wr.result = WorkRequest.Results.SUCCESS
        wr.save()

        msg = r"^Only aborted or failed tasks can be retried$"

        self.assertFalse(wr.can_retry())
        with self.assertRaisesRegex(CannotRetry, msg):
            wr.retry()

        wr.status = WorkRequest.Statuses.COMPLETED
        wr.result = WorkRequest.Results.ERROR
        wr.save()
        self.assertFalse(wr.can_retry())
        with self.assertRaisesRegex(CannotRetry, msg):
            wr.retry()

        wr.status = WorkRequest.Statuses.PENDING
        wr.result = WorkRequest.Results.NONE
        wr.save()
        self.assertFalse(wr.can_retry())
        with self.assertRaisesRegex(CannotRetry, msg):
            wr.retry()

        wr.status = WorkRequest.Statuses.RUNNING
        wr.result = WorkRequest.Results.NONE
        wr.save()
        self.assertFalse(wr.can_retry())
        with self.assertRaisesRegex(CannotRetry, msg):
            wr.retry()

        wr.status = WorkRequest.Statuses.BLOCKED
        wr.result = WorkRequest.Results.NONE
        wr.save()
        self.assertFalse(wr.can_retry())
        with self.assertRaisesRegex(CannotRetry, msg):
            wr.retry()

    def test_can_retry_task_type(self) -> None:
        """Only some task types can be retried."""
        wr = self.workflow.create_child("noop")
        wr.status = WorkRequest.Statuses.COMPLETED
        wr.result = WorkRequest.Results.SUCCESS
        wr.save()

        msg = (
            r"^Only workflow, worker, internal, or signing tasks"
            " can be retried$"
        )

        wr.task_type = TaskTypes.SERVER
        wr.status = WorkRequest.Statuses.ABORTED
        wr.result = WorkRequest.Results.NONE
        wr.save()
        self.assertFalse(wr.can_retry())
        with self.assertRaisesRegex(CannotRetry, msg):
            wr.retry()

    def test_cant_retry_invalid_task_data(self) -> None:
        """Work request with invalid task data cannot be retried."""
        wr = self.workflow.create_child("noop")
        wr.status = WorkRequest.Statuses.ABORTED
        wr.result = WorkRequest.Results.NONE
        wr.task_name = "sbuild"
        wr.save()
        self.assertFalse(wr.can_retry())
        msg = "^Task dependencies cannot be satisfied$"
        with self.assertRaisesRegex(CannotRetry, msg):
            wr.retry()

    def test_cant_retry_if_lookups_fail(self) -> None:
        """Work requests with unsatisfiable lookups can't retry."""
        wr = self.workflow.create_child("noop")
        wr.status = WorkRequest.Statuses.ABORTED
        wr.result = WorkRequest.Results.NONE
        wr.save()

        collection = self.create_collection(
            name="debian",
            category=CollectionCategory.ENVIRONMENTS,
        )

        with preserve_task_registry():

            class LookupTask(Noop):
                """Test task that performs a lookup."""

                def compute_dynamic_data(
                    self, task_database: TaskDatabaseInterface
                ) -> BaseDynamicTaskData | None:
                    """Perform a lookup."""
                    task_database.lookup_single_artifact(
                        "debian@debian:environments/match:codename=bookworm"
                    )
                    return None

            wr = self.create_work_request(
                task_name="lookuptask",
                task_data={},
                status=WorkRequest.Statuses.ABORTED,
            )

            self.assertFalse(wr.can_retry())
            msg = r"^Task dependencies cannot be satisfied$"
            with self.assertRaisesRegex(CannotRetry, msg):
                wr.retry()

            # Create the looked up artifact
            tarball, _ = self.create_artifact(
                category=ArtifactCategory.SYSTEM_TARBALL,
                data={"architecture": "amd64", "codename": "bookworm"},
            )
            collection.manager.add_artifact(tarball, user=self.user)

            self.assertTrue(wr.can_retry())
            wr.retry()

    def test_retry_supersede(self) -> None:
        """Test behaviour of retry, superseding work requests."""
        # An aborted workflow, with a failed task and one depending on it
        self.workflow.status = WorkRequest.Statuses.ABORTED
        self.workflow.result = WorkRequest.Results.FAILURE
        self.workflow.completed_at = timezone.now()
        self.workflow.save()

        wr_failed = self.workflow.create_child("noop")
        wr_failed.status = WorkRequest.Statuses.COMPLETED
        wr_failed.result = WorkRequest.Results.FAILURE
        wr_failed.save()

        # Two reverse dependencies
        wr_dep1 = self.workflow.create_child(
            "noop", status=WorkRequest.Statuses.PENDING
        )
        wr_dep1.dependencies.add(wr_failed)
        wr_dep2 = self.workflow.create_child(
            "noop", status=WorkRequest.Statuses.COMPLETED
        )
        wr_dep2.dependencies.add(wr_failed)
        # One forward dependency
        wr_fdep1 = self.workflow.create_child(
            "noop", status=WorkRequest.Statuses.COMPLETED
        )
        wr_failed.dependencies.add(wr_fdep1)

        # wr_failed gets retried
        self.assertTrue(wr_failed.can_retry())
        wr_new = wr_failed.retry()

        self.assertEqual(wr_new.workspace, wr_failed.workspace)
        self.assertGreater(wr_new.created_at, wr_failed.created_at)
        self.assertIsNone(wr_new.started_at)
        self.assertIsNone(wr_new.completed_at)
        self.assertEqual(wr_new.created_by, wr_failed.created_by)
        self.assertEqual(wr_new.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(wr_new.result, WorkRequest.Results.NONE)
        self.assertIsNone(wr_new.worker)
        self.assertEqual(wr_new.task_type, wr_failed.task_type)
        self.assertEqual(wr_new.task_name, wr_failed.task_name)
        self.assertEqual(wr_new.task_data, wr_failed.task_data)
        self.assertIsNone(wr_new.dynamic_task_data)
        self.assertEqual(wr_new.priority_base, wr_failed.priority_base)
        self.assertEqual(
            wr_new.priority_adjustment, wr_failed.priority_adjustment
        )
        self.assertEqual(wr_new.unblock_strategy, wr_failed.unblock_strategy)
        self.assertEqual(wr_new.parent, self.workflow)
        self.assertEqual(
            wr_new.workflow_data_json, wr_failed.workflow_data_json
        )
        self.assertEqual(
            wr_new.event_reactions_json, wr_failed.event_reactions_json
        )
        self.assertIsNone(wr_new.internal_collection)
        self.assertEqual(wr_new.expiration_delay, wr_failed.expiration_delay)
        self.assertEqual(wr_new.supersedes, wr_failed)
        self.assertEqual(wr_failed.superseded, wr_new)

        self.assertQuerysetEqual(wr_new.dependencies.all(), [wr_fdep1])

        self.assertQuerysetEqual(wr_dep1.dependencies.all(), [wr_new])
        wr_dep1.refresh_from_db()
        self.assertEqual(wr_dep1.status, WorkRequest.Statuses.BLOCKED)

        self.assertQuerysetEqual(wr_dep2.dependencies.all(), [wr_new])
        wr_dep2.refresh_from_db()
        self.assertEqual(wr_dep2.status, WorkRequest.Statuses.COMPLETED)

        self.workflow.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(self.workflow.result, WorkRequest.Results.NONE)
        self.assertIsNone(self.workflow.completed_at)

    def test_retry_in_running_workflow(self) -> None:
        """Test retrying a task in a running workflow."""
        # An aborted workflow, with a failed task and one depending on it
        now = timezone.now()
        self.workflow.status = WorkRequest.Statuses.RUNNING
        # Bogus value to check that it does not get reset
        self.workflow.completed_at = now
        self.workflow.save()

        wr_failed = self.workflow.create_child("noop")
        wr_failed.status = WorkRequest.Statuses.COMPLETED
        wr_failed.result = WorkRequest.Results.FAILURE
        wr_failed.save()

        # wr_failed gets retried
        wr_failed.retry()

        self.workflow.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(self.workflow.completed_at, now)

    def test_retry_in_failed_workflow(self) -> None:
        """Test retrying a task in a failed workflow."""
        # An aborted workflow, with a failed task and one depending on it
        now = timezone.now()
        self.workflow.status = WorkRequest.Statuses.COMPLETED
        self.workflow.result = WorkRequest.Results.FAILURE
        # Bogus value to check that it does not get reset
        self.workflow.completed_at = now
        self.workflow.save()

        wr_failed = self.workflow.create_child("noop")
        wr_failed.status = WorkRequest.Statuses.COMPLETED
        wr_failed.result = WorkRequest.Results.FAILURE
        wr_failed.save()

        # wr_failed gets retried
        wr_failed.retry()

        # The workflow is set to running again
        self.workflow.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.RUNNING)
        self.assertIsNone(self.workflow.completed_at)

    def test_retry_empty_workflow(self) -> None:
        """Test behaviour of retrying for an empty workflow."""
        self.workflow.status = WorkRequest.Statuses.ABORTED
        self.workflow.save()

        wr = self.workflow.retry()

        # Retrying a workflow does not create a new work request
        self.assertEqual(wr.pk, self.workflow.pk)

        self.workflow.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(self.workflow.result, WorkRequest.Results.SUCCESS)

    def test_retry_workflow(self) -> None:
        """Test behaviour of retrying for an empty workflow."""
        wr_completed = self.workflow.create_child("noop")
        wr_aborted = self.workflow.create_child("noop")

        # Simulate the previous workflow lifecycle
        self.assertTrue(self.workflow.mark_running())
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.RUNNING)
        self.workflow.unblock_workflow_children()
        wr_completed.refresh_from_db()
        self.assertTrue(
            wr_completed.mark_completed(WorkRequest.Results.SUCCESS)
        )
        wr_completed.completed_at = (
            orig_completed := timezone.now() - timedelta(days=7)
        )
        wr_completed.save()
        wr_aborted.refresh_from_db()
        self.assertTrue(wr_aborted.mark_aborted())

        # The workflow ended up as aborted
        self.workflow.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.ABORTED)

        # Retry it
        wr = self.workflow.retry()
        self.assertEqual(wr.pk, self.workflow.pk)

        self.workflow.refresh_from_db()
        wr_completed.refresh_from_db()
        wr_aborted.refresh_from_db()

        # The workflow is running
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(self.workflow.result, WorkRequest.Results.NONE)
        self.assertEqual(wr_completed.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr_completed.result, WorkRequest.Results.SUCCESS)
        # Successful work requests did not get restarted
        self.assertEqual(wr_completed.completed_at, orig_completed)
        self.assertEqual(wr_aborted.status, WorkRequest.Statuses.ABORTED)
        self.assertEqual(wr_aborted.result, WorkRequest.Results.NONE)
        wr_aborted_new = wr_aborted.superseded
        self.assertEqual(wr_aborted_new.status, WorkRequest.Statuses.PENDING)
        self.assertEqual(wr_aborted_new.result, WorkRequest.Results.NONE)

        # Complete the one running task
        wr_aborted_new.mark_completed(WorkRequest.Results.SUCCESS)
        self.workflow.refresh_from_db()

        # Workflow completed successfully
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(self.workflow.result, WorkRequest.Results.SUCCESS)

    def test_retry_workflow_calls(self) -> None:
        """Test behaviour of retrying for an empty workflow."""
        wr_aborted = self.workflow.create_child("noop")

        # Simulate the previous workflow lifecycle
        self.assertTrue(self.workflow.mark_running())
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.RUNNING)
        self.workflow.unblock_workflow_children()
        wr_aborted.refresh_from_db()
        self.assertTrue(wr_aborted.mark_aborted())

        # The workflow ended up as aborted
        self.workflow.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.ABORTED)

        with (
            mock.patch(
                "debusine.server.scheduler.populate_workflow"
            ) as populate_workflow,
            mock.patch(
                "debusine.db.models.work_requests.WorkRequest"
                ".maybe_finish_workflow"
            ) as maybe_finish_workflow,
        ):
            self.workflow.retry()

        populate_workflow.assert_called_with(self.workflow)
        maybe_finish_workflow.assert_called()

    def test_retry_workflow_child_fails(self) -> None:
        """Test retrying a workflow with a child that cannot retry."""
        wr_aborted1 = self.workflow.create_child("noop")
        wr_unretriable = self.workflow.create_child("noop")
        wr_aborted2 = self.workflow.create_child("noop")

        self.assertTrue(self.workflow.mark_running())
        self.workflow.unblock_workflow_children()
        wr_aborted1.refresh_from_db()
        wr_unretriable.refresh_from_db()
        wr_aborted2.refresh_from_db()
        self.assertTrue(wr_aborted1.mark_aborted())
        self.assertTrue(wr_unretriable.mark_aborted())
        self.assertTrue(wr_aborted2.mark_aborted())

        # The workflow ended up as aborted
        self.workflow.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.ABORTED)

        # Make wr_unretriable unretriable
        orig_check_retry = WorkRequest._check_retry

        def test_check_retry(self):
            if self.pk == wr_unretriable.pk:
                raise CannotRetry("expected failure")
            else:
                orig_check_retry(self)

        try:
            with transaction.atomic():
                with (
                    mock.patch(
                        "debusine.db.models.WorkRequest._check_retry",
                        side_effect=test_check_retry,
                        autospec=True,
                    ),
                    self.assertRaises(CannotRetry) as exc,
                ):
                    self.workflow.retry()

                self.assertEqual(
                    exc.exception.args[0],
                    f"child work request {wr_unretriable.pk} cannot be retried",
                )
                self.assertIsNotNone(exc.exception.__context__)
                assert exc.exception.__context__ is not None
                self.assertEqual(
                    exc.exception.__context__.args[0], "expected failure"
                )

                self.workflow.refresh_from_db()
                wr_aborted1.refresh_from_db()
                wr_aborted1_new = wr_aborted1.superseded
                wr_unretriable.refresh_from_db()
                wr_aborted2.refresh_from_db()
                self.assertFalse(hasattr(wr_aborted2, "superseded"))
                self.assertEqual(
                    self.workflow.status, WorkRequest.Statuses.RUNNING
                )
                self.assertEqual(
                    wr_aborted1_new.status, WorkRequest.Statuses.PENDING
                )
                self.assertEqual(
                    wr_unretriable.status, WorkRequest.Statuses.ABORTED
                )
                self.assertEqual(
                    wr_aborted2.status, WorkRequest.Statuses.ABORTED
                )
                # Raise to roll back the transaction
                raise exc.exception
        except CannotRetry as exc1:
            self.assertIs(exc1, exc.exception)

        self.workflow.refresh_from_db()
        wr_aborted1.refresh_from_db()
        wr_unretriable.refresh_from_db()
        wr_aborted2.refresh_from_db()
        self.assertEqual(self.workflow.status, WorkRequest.Statuses.ABORTED)
        self.assertEqual(wr_aborted1.status, WorkRequest.Statuses.ABORTED)
        self.assertEqual(wr_unretriable.status, WorkRequest.Statuses.ABORTED)
        self.assertEqual(wr_aborted2.status, WorkRequest.Statuses.ABORTED)

    def test_is_part_of_workflow_standalone(self) -> None:
        """A standalone work request is not part of a workflow."""
        wr = self.create_work_request()

        self.assertFalse(wr.is_part_of_workflow)

    def test_is_part_of_workflow_root(self) -> None:
        """A work request of type WORKFLOW is part of a workflow."""
        self.assertTrue(self.workflow.is_part_of_workflow)

    def test_is_part_of_workflow_child(self) -> None:
        """A work request with a parent is part of a workflow."""
        wr = self.create_work_request(parent=self.workflow)

        self.assertTrue(wr.is_part_of_workflow)

    def test_get_workflow_root_standalone(self) -> None:
        """A standalone work request has no workflow root."""
        wr = self.create_work_request()

        self.assertIsNone(wr.get_workflow_root())

    def test_get_workflow_root_root(self) -> None:
        """A top-level work request in a workflow is its own root."""
        self.assertEqual(self.workflow.get_workflow_root(), self.workflow)

    def test_get_workflow_root_child(self) -> None:
        """A child work request in a workflow has a root."""
        wr = self.create_work_request(parent=self.workflow)

        self.assertEqual(wr.get_workflow_root(), self.workflow)

    def test_get_workflow_root_grandchild(self) -> None:
        """A grandchild work request in a workflow has a root."""
        root = self.workflow
        child_template = WorkflowTemplate.objects.create(
            name="child",
            workspace=default_workspace(),
            task_name="noop",
            task_data={},
        )
        child = WorkRequest.objects.create_workflow(
            template=child_template,
            data={},
            created_by=self.get_test_user(),
            parent=root,
        )
        wr = self.create_work_request(parent=child)

        self.assertEqual(wr.get_workflow_root(), root)

    def test_workflow_display_name_explicit(self) -> None:
        """`workflow_display_name` returns an explicitly-set display name."""
        wr = self.create_work_request(
            workflow_data_json={"display_name": "Foo", "step": "foo"}
        )

        self.assertEqual(wr.workflow_display_name, "Foo")

    def test_workflow_display_name_task_name(self) -> None:
        """`workflow_display_name` defaults to the task name."""
        wr = self.create_work_request(task_name="noop")

        self.assertEqual(wr.workflow_display_name, "noop")

    def test_effective_expiration_delay(self):
        """Test WorkRequest.effective_expiration_delay() method."""
        wr = self.create_work_request(task_name="noop")

        wr.expiration_delay = timedelta(days=1)
        self.assertEqual(wr.effective_expiration_delay(), timedelta(days=1))

        wr.expiration_delay = timedelta(days=0)
        self.assertEqual(wr.effective_expiration_delay(), timedelta(days=0))

        workspace = Workspace.objects.create_with_name(name="test")
        workspace.default_expiration_delay = timedelta(days=2)
        wr.workspace = workspace

        wr.expiration_delay = timedelta(days=1)
        self.assertEqual(wr.effective_expiration_delay(), timedelta(days=1))

        wr.expiration_delay = timedelta(days=0)
        self.assertEqual(wr.effective_expiration_delay(), timedelta(days=0))

        wr.expiration_delay = None
        self.assertEqual(wr.effective_expiration_delay(), timedelta(days=2))


class WorkRequestWorkerTests(ChannelsHelpersMixin, TestHelpersMixin, TestCase):
    """Tests for the WorkRequest class scheduling."""

    def setUp(self):
        """Set up WorkRequest to be used in the tests."""
        worker = Worker.objects.create_with_fqdn(
            "computer.lan", Token.objects.create()
        )

        self.work_request = self.create_work_request(
            task_name='request-01',
            worker=worker,
            status=WorkRequest.Statuses.PENDING,
        )

    def test_str(self):
        """Test WorkerRequest.__str__ return WorkRequest.id."""
        self.assertEqual(self.work_request.__str__(), str(self.work_request.id))

    def test_mark_running_from_aborted(self):
        """Test WorkRequest.mark_running() doesn't change (was aborted)."""
        self.work_request.status = WorkRequest.Statuses.ABORTED
        self.work_request.save()

        self.assertFalse(self.work_request.mark_running())

        self.work_request.refresh_from_db()
        self.assertEqual(self.work_request.status, WorkRequest.Statuses.ABORTED)

    def test_mark_running(self):
        """Test WorkRequest.mark_running() change status to running."""
        self.work_request.status = WorkRequest.Statuses.PENDING
        self.work_request.save()
        self.assertIsNone(self.work_request.started_at)

        self.assertTrue(self.work_request.mark_running())

        self.work_request.refresh_from_db()
        self.assertEqual(self.work_request.status, WorkRequest.Statuses.RUNNING)
        self.assertLess(self.work_request.started_at, timezone.now())

        # Marking as running again (a running WorkRequest) is a no-op
        started_at = self.work_request.started_at
        self.assertTrue(self.work_request.mark_running())

        self.work_request.refresh_from_db()
        self.assertEqual(self.work_request.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(self.work_request.started_at, started_at)

    def test_mark_running_fails_worker_already_running(self):
        """WorkRequest.mark_running() return False: worker already running."""
        self.work_request.status = WorkRequest.Statuses.PENDING
        self.work_request.save()

        other_work_request = self.create_work_request(
            status=WorkRequest.Statuses.RUNNING, worker=self.work_request.worker
        )

        with self.assertLogsContains(
            f"Cannot mark WorkRequest {self.work_request.pk} as running - the "
            f"assigned worker {self.work_request.worker} is running too many "
            f"other WorkRequests: [{other_work_request!r}]",
            logger="debusine.db.models",
            level=logging.DEBUG,
        ):
            self.assertFalse(self.work_request.mark_running())

    def test_mark_running_fails_no_assigned_worker(self):
        """WorkRequest.mark_running() return False: no worker assigned."""
        self.work_request.status = WorkRequest.Statuses.PENDING
        self.work_request.worker = None
        self.work_request.save()

        self.assertFalse(self.work_request.mark_running())

    def test_mark_completed_from_aborted(self):
        """Test WorkRequest.mark_completed() doesn't change (was aborted)."""
        self.work_request.status = WorkRequest.Statuses.ABORTED
        self.work_request.save()

        self.work_request.refresh_from_db()

        self.assertFalse(
            self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        )

        self.work_request.refresh_from_db()
        self.assertEqual(self.work_request.status, WorkRequest.Statuses.ABORTED)

    def test_mark_completed(self):
        """Test WorkRequest.mark_completed() changes status to completed."""
        self.work_request.status = WorkRequest.Statuses.RUNNING
        self.work_request.save()

        self.assertIsNone(self.work_request.completed_at)
        self.assertEqual(self.work_request.result, WorkRequest.Results.NONE)

        self.assertTrue(
            self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        )

        self.work_request.refresh_from_db()
        self.assertEqual(
            self.work_request.status, WorkRequest.Statuses.COMPLETED
        )
        self.assertEqual(self.work_request.result, WorkRequest.Results.SUCCESS)
        self.assertLess(self.work_request.completed_at, timezone.now())

    def test_mark_aborted(self):
        """Test WorkRequest.mark_aborted() changes status to aborted."""
        self.assertIsNone(self.work_request.completed_at)

        self.assertTrue(self.work_request.mark_aborted())

        self.work_request.refresh_from_db()
        self.assertEqual(self.work_request.status, WorkRequest.Statuses.ABORTED)
        self.assertLess(self.work_request.completed_at, timezone.now())

    def test_maybe_finish_workflow_not_workflow(self):
        """`maybe_finish_workflow` does nothing on non-workflows."""
        self.assertFalse(self.work_request.maybe_finish_workflow())

    def test_maybe_finish_workflow_already_aborted(self):
        """`maybe_finish_workflow` does nothing if it was already aborted."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        workflow.create_child("noop", status=WorkRequest.Statuses.COMPLETED)
        self.assertTrue(workflow.mark_aborted())

        self.assertFalse(workflow.maybe_finish_workflow())

        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.ABORTED)

    def test_maybe_finish_workflow_has_children_in_progress(self):
        """`maybe_finish_workflow` does nothing if a child is in progress."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        children = [workflow.create_child("noop") for _ in range(2)]
        workflow.mark_running()
        workflow.unblock_workflow_children()
        children[0].refresh_from_db()
        self.assertTrue(children[0].mark_completed(WorkRequest.Results.SUCCESS))

        self.assertFalse(workflow.maybe_finish_workflow())

        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.RUNNING)

    def test_maybe_finish_workflow_all_succeeded(self):
        """All children succeeded: COMPLETED/SUCCESS."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        children = [workflow.create_child("noop") for _ in range(2)]
        workflow.mark_running()
        workflow.unblock_workflow_children()

        for child in children:
            child.refresh_from_db()
            self.assertTrue(child.mark_completed(WorkRequest.Results.SUCCESS))

        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(workflow.result, WorkRequest.Results.SUCCESS)
        self.assertLess(workflow.completed_at, timezone.now())

    def test_maybe_finish_workflow_allow_failure(self):
        """Some children failed but are allowed to fail: COMPLETED/SUCCESS."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        children = [
            workflow.create_child(
                "noop",
                workflow_data=WorkRequestWorkflowData(
                    allow_failure=True, display_name="test", step="test"
                ),
            )
            for _ in range(2)
        ]
        workflow.mark_running()
        workflow.unblock_workflow_children()

        children[0].refresh_from_db()
        children[0].mark_completed(WorkRequest.Results.FAILURE)
        children[1].refresh_from_db()
        children[1].mark_completed(WorkRequest.Results.SUCCESS)

        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(workflow.result, WorkRequest.Results.SUCCESS)
        self.assertLess(workflow.completed_at, timezone.now())

    def test_maybe_finish_workflow_some_failed(self):
        """Some children failed: COMPLETED/FAILURE."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        children = [workflow.create_child("noop") for _ in range(2)]
        workflow.mark_running()
        workflow.unblock_workflow_children()

        children[0].refresh_from_db()
        children[0].mark_completed(WorkRequest.Results.FAILURE)
        children[1].refresh_from_db()
        children[1].mark_completed(WorkRequest.Results.SUCCESS)

        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(workflow.result, WorkRequest.Results.FAILURE)
        self.assertLess(workflow.completed_at, timezone.now())

    def test_maybe_finish_workflow_some_aborted(self):
        """Some children were aborted: `maybe_finish_workflow` sets ABORTED."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        children = [workflow.create_child("noop") for _ in range(2)]
        workflow.mark_running()
        workflow.unblock_workflow_children()

        children[0].refresh_from_db()
        children[0].mark_aborted()
        children[1].refresh_from_db()
        children[1].mark_completed(WorkRequest.Results.SUCCESS)

        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.ABORTED)
        self.assertLess(workflow.completed_at, timezone.now())

    def test_maybe_finish_workflow_superseded_aborted(self):
        """Some aborted children were superseded: ignores them."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        children = [workflow.create_child("noop") for _ in range(2)]
        workflow.started_at = timezone.now()
        workflow.status = WorkRequest.Statuses.RUNNING
        workflow.save()

        children[0].completed_at = timezone.now()
        children[0].status = WorkRequest.Statuses.ABORTED
        children[0].save()
        children[1].completed_at = timezone.now()
        children[1].status = WorkRequest.Statuses.COMPLETED
        children[1].result = WorkRequest.Results.SUCCESS
        children[1].save()

        self.assertTrue(workflow.maybe_finish_workflow())
        self.assertEqual(workflow.status, WorkRequest.Statuses.ABORTED)
        self.assertLess(workflow.completed_at, timezone.now())
        prev_completed_at = workflow.completed_at

        new_wr = workflow.create_child("noop")
        new_wr.supersedes = children[0]
        new_wr.status = WorkRequest.Statuses.COMPLETED
        new_wr.result = WorkRequest.Results.SUCCESS
        new_wr.save()
        workflow.status = WorkRequest.Statuses.RUNNING
        workflow.save()

        self.assertTrue(workflow.maybe_finish_workflow())
        self.assertEqual(workflow.status, WorkRequest.Statuses.COMPLETED)
        self.assertLess(workflow.completed_at, timezone.now())
        self.assertGreater(workflow.completed_at, prev_completed_at)

    async def test_assign_worker(self):
        """Assign Worker to WorkRequest."""

        def setup_work_request() -> Worker:
            worker = self.work_request.worker

            self.work_request.worker = None
            self.work_request.save()

            self.work_request.refresh_from_db()

            # Initial status (no worker)
            self.assertIsNone(self.work_request.worker)
            self.assertEqual(
                self.work_request.status, WorkRequest.Statuses.PENDING
            )

            return worker

        worker = await sync_to_async(setup_work_request)()

        channel = await self.create_channel(worker.token.hash)

        # Assign the worker to the WorkRequest
        await sync_to_async(self.work_request.assign_worker)(worker)

        await self.assert_channel_received(
            channel, {"type": "work_request.assigned"}
        )

        def assert_final_status():
            self.work_request.refresh_from_db()

            # Assert final status
            self.assertEqual(self.work_request.worker, worker)
            self.assertEqual(
                self.work_request.status, WorkRequest.Statuses.PENDING
            )

        await sync_to_async(assert_final_status)()

    def test_duration(self):
        """duration() returns the correct duration."""
        self.work_request.started_at = datetime(2022, 3, 7, 10, 51)
        self.work_request.completed_at = datetime(2022, 3, 9, 12, 53)
        duration = self.work_request.completed_at - self.work_request.started_at

        self.assertEqual(
            self.work_request.duration,
            duration.total_seconds(),
        )

    def test_duration_is_none_completed_at_is_none(self):
        """duration() returns None because completed_at is None."""
        self.work_request.started_at = datetime(2022, 3, 7, 10, 51)
        self.work_request.completed_at = None

        self.assertIsNone(self.work_request.duration)

    def test_duration_is_none_started_at_is_none(self):
        """duration() returns None because started_at is None."""
        self.work_request.started_at = None
        self.work_request.completed_at = datetime(2022, 3, 7, 10, 51)

        self.assertIsNone(self.work_request.duration)

    def test_priority_effective(self):
        """priority_effective() returns the effective priority."""
        self.work_request.priority_base = 10
        self.work_request.priority_adjustment = 1

        self.assertEqual(self.work_request.priority_effective, 11)

    def test_invalid_pending(self):
        """Reject invalid pending state change."""
        self.work_request.status = WorkRequest.Statuses.RUNNING
        self.assertFalse(self.work_request.mark_pending())
        self.assertEqual(self.work_request.status, WorkRequest.Statuses.RUNNING)

    def test_add_dependency_within_workflow(self):
        """Dependencies within a workflow are allowed."""
        template = WorkflowTemplate.objects.create(
            name="test", workspace=default_workspace(), task_name="noop"
        )
        workflow = WorkRequest.objects.create_workflow(
            template=template, created_by=self.get_test_user(), data={}
        )
        children = [workflow.create_child("noop") for _ in range(2)]

        children[0].add_dependency(children[1])

        self.assertQuerysetEqual(children[0].dependencies.all(), [children[1]])

    def test_add_dependency_outside_workflow(self):
        """Dependencies from within a workflow to outside are forbidden."""
        template = WorkflowTemplate.objects.create(
            name="test", workspace=default_workspace(), task_name="noop"
        )
        workflow = WorkRequest.objects.create_workflow(
            template=template, created_by=self.get_test_user(), data={}
        )
        child = workflow.create_child("noop")
        outside = self.create_work_request()

        with self.assertRaisesRegex(
            ValueError,
            "Work requests in a workflow may not depend on other work "
            "requests outside that workflow",
        ):
            child.add_dependency(outside)

    def test_add_dependency_auto_block(self):
        """`add_dependency` automatically blocks PENDING/DEPS work requests."""
        rdep = self.create_work_request()
        rdep.add_dependency(self.work_request)
        self.assertEqual(rdep.status, WorkRequest.Statuses.BLOCKED)

    def test_add_dependency_no_auto_block_not_pending(self):
        """`add_dependency` doesn't automatically block non-PENDING requests."""
        rdep = self.create_work_request()
        rdep.status = WorkRequest.Statuses.ABORTED
        rdep.save()
        rdep.add_dependency(self.work_request)
        self.assertEqual(rdep.status, WorkRequest.Statuses.ABORTED)

    def test_add_dependency_no_auto_block_not_deps(self):
        """`add_dependency` doesn't automatically block non-DEPS requests."""
        rdep = self.create_work_request()
        rdep.unblock_strategy = WorkRequest.UnblockStrategy.MANUAL
        rdep.save()
        rdep.add_dependency(self.work_request)
        self.assertEqual(rdep.status, WorkRequest.Statuses.PENDING)

    def test_add_dependency_no_auto_block_already_completed(self):
        """`add_dependency` doesn't auto-block on completed dependencies."""
        self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        rdep = self.create_work_request()
        rdep.add_dependency(self.work_request)
        self.assertEqual(rdep.status, WorkRequest.Statuses.PENDING)

    def create_rdep(self) -> WorkRequest:
        """Create a reverse-dependency for the default work request."""
        rdep = self.create_work_request(
            task_name='rdep', worker=self.work_request.worker
        )
        rdep.add_dependency(self.work_request)
        return rdep

    def test_dependencies_failure(self):
        """Failure result gets rdep aborted."""
        rdep = self.create_rdep()
        self.work_request.mark_completed(WorkRequest.Results.FAILURE)
        rdep.refresh_from_db()
        self.assertEqual(rdep.status, WorkRequest.Statuses.ABORTED)

    def test_dependencies_success_only_dep(self):
        """Success result, with nothing else to do, gets rdep pending."""
        rdep = self.create_rdep()
        self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        rdep.refresh_from_db()
        self.assertEqual(rdep.status, WorkRequest.Statuses.PENDING)

    def test_dependencies_allow_failure(self):
        """Failure result, with allow_failure, gets rdep pending."""
        rdep = self.create_rdep()
        self.work_request.workflow_data = WorkRequestWorkflowData(
            allow_failure=True, display_name="test", step="test"
        )
        self.work_request.mark_completed(WorkRequest.Results.FAILURE)
        rdep.refresh_from_db()
        self.assertEqual(rdep.status, WorkRequest.Statuses.PENDING)

    def test_dependencies_success_manual(self):
        """Success result with manual unblock keeps rdep blocked."""
        rdep = self.create_rdep()
        rdep.unblock_strategy = WorkRequest.UnblockStrategy.MANUAL
        rdep.save()
        self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        rdep.refresh_from_db()
        self.assertEqual(rdep.status, WorkRequest.Statuses.BLOCKED)

    def test_dependencies_success_more_deps(self):
        """Success result with something else to do keeps rdep blocked."""
        rdep = self.create_rdep()
        work_request2 = self.create_work_request(
            task_name='request-02',
            worker=self.work_request.worker,
            status=WorkRequest.Statuses.RUNNING,
        )
        rdep.add_dependency(work_request2)
        self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        rdep.refresh_from_db()
        self.assertEqual(rdep.status, WorkRequest.Statuses.BLOCKED)

    def test_dependencies_in_pending_workflow(self):
        """Parent workflow that isn't running yet keeps rdep blocked."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        rdep = workflow.create_child("noop")
        self.work_request.parent = workflow
        self.work_request.save()
        rdep.add_dependency(self.work_request)
        self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        rdep.refresh_from_db()
        self.assertEqual(rdep.status, WorkRequest.Statuses.BLOCKED)

    def test_dependencies_in_running_workflow(self):
        """Parent workflow in running status gets rdep pending."""
        template = self.create_workflow_template("test", "noop")
        workflow = WorkRequest.objects.create_workflow(
            template=template, data={}, created_by=self.get_test_user()
        )
        rdep = workflow.create_child("noop")
        self.work_request.parent = workflow
        self.work_request.save()
        rdep.add_dependency(self.work_request)
        workflow.mark_running()
        self.work_request.mark_completed(WorkRequest.Results.SUCCESS)
        rdep.refresh_from_db()
        self.assertEqual(rdep.status, WorkRequest.Statuses.PENDING)

    def test_create_child_non_workflow(self):
        """Non-WORKFLOW work requests may not have children."""
        with self.assertRaisesRegex(
            ValueError, r"Only workflows may have child work requests\."
        ):
            self.work_request.create_child("noop")

    def test_event_reactions_update_collection(self):
        """Update collection on task completion."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.get_test_user(),
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
        )
        c = self.create_collection(
            name=f"workflow-{wr.id}",
            category=CollectionCategory.WORKFLOW_INTERNAL,
        )
        cm = CollectionManagerInterface.get_manager_for(c)
        a_replaced, _ = self.create_artifact(
            category=ArtifactCategory.TEST,
            data={
                "deb_fields": {
                    "Package": "actiontest",
                    "Version": "1.0-1",
                },
                "type": "rebuild",
            },
        )
        cm.add_artifact(
            a_replaced, user=wr.created_by, name="actiontest_rebuild"
        )
        item = cm.lookup("name:actiontest_rebuild")
        self.assertEqual(item.artifact, a_replaced)

        a_ignored, _ = self.create_artifact(
            category=ArtifactCategory.TEST,
            data={
                "deb_fields": {
                    "Package": "ignored",
                    "Version": "1.0-1",
                },
                "type": "rebuild",
            },
        )
        cm.add_artifact(a_ignored, user=wr.created_by, name="ignored_rebuild")
        item = cm.lookup("name:ignored_rebuild")
        self.assertEqual(item.artifact, a_ignored)

        a_replacement, _ = self.create_artifact(
            category=ArtifactCategory.TEST,
            data={
                'deb_fields': {
                    "Package": "actiontest",
                    "Version": "1.0-1+b1",
                },
                "type": "rebuild",
            },
            work_request=wr,
        )

        a_added, _ = self.create_artifact(
            category=ArtifactCategory.TEST,
            data={
                "deb_fields": {
                    "Package": "actiontest2",
                    "Version": "2.0-1+b1",
                },
                "type": "rebuild",
            },
            work_request=wr,
        )

        action = ActionUpdateCollectionWithArtifacts.parse_obj(
            {
                "artifact_filters": {
                    "category": ArtifactCategory.TEST,
                    "data__deb_fields__Package__startswith": "actiontest",
                },
                "collection": "internal@collections",
                "name_template": "{name}_{type}",
                "variables": {
                    "$name": "deb_fields.Package",
                    "$type": "type",
                },
            }
        )
        wr.event_reactions = EventReactions(
            on_success=[action],
            on_failure=[action],
        )
        wr.mark_completed(WorkRequest.Results.SUCCESS)
        self.assertIsNotNone(
            CollectionItem.objects.get(
                parent_collection=c, artifact=a_replaced
            ).removed_at
        )
        item = cm.lookup("name:ignored_rebuild")
        self.assertEqual(item.artifact, a_ignored)
        item = cm.lookup("name:actiontest_rebuild")
        self.assertEqual(item.artifact, a_replacement)
        item = cm.lookup("name:actiontest2_rebuild")
        self.assertEqual(item.artifact.id, a_added.id)

        # For completeness, test NONE
        wr.result = WorkRequest.Results.NONE
        actions = wr.get_triggered_actions()
        for action_type, actions in actions.items():
            self.assertEqual(len(actions), 0)

    def test_event_reactions_update_collection_no_name_template(self):
        """Without name_template, variables are passed through."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.get_test_user(),
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
        )
        c = Collection.objects.create(
            name="debian",
            category=CollectionCategory.ENVIRONMENTS,
            workspace=default_workspace(),
        )
        tarball, _ = self.create_artifact(
            category=ArtifactCategory.SYSTEM_TARBALL,
            data={"architecture": "amd64"},
            work_request=wr,
        )
        action = ActionUpdateCollectionWithArtifacts.parse_obj(
            {
                "artifact_filters": {
                    "category": ArtifactCategory.SYSTEM_TARBALL
                },
                "collection": "debian@debian:environments",
                "variables": {"codename": "bookworm", "variant": "buildd"},
            }
        )
        wr.event_reactions = EventReactions(on_success=[action])

        wr.mark_completed(WorkRequest.Results.SUCCESS)

        self.assertIsNotNone(
            c.manager.lookup(
                "match:format=tarball:codename=bookworm:architecture=amd64:"
                "variant=buildd"
            ).artifact,
            tarball,
        )

    def test_event_reactions_update_collection_constraint_violation(self):
        """Violating collection constraints is an error."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.get_test_user(),
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
        )
        self.create_collection(
            name=f"workflow-{wr.id}",
            category=CollectionCategory.WORKFLOW_INTERNAL,
        )
        self.create_artifact(category=ArtifactCategory.TEST, work_request=wr)
        # This action does not specify an item name.
        action = ActionUpdateCollectionWithArtifacts(
            artifact_filters={"category": ArtifactCategory.TEST},
            collection="internal@collections",
        )
        wr.event_reactions = EventReactions(on_success=[action])
        wr.status = WorkRequest.Statuses.RUNNING
        with self.assertLogsContains(
            "Cannot replace or add artifact",
            logger="debusine.db.models",
            level=logging.ERROR,
        ):
            wr.mark_completed(WorkRequest.Results.SUCCESS)

    def test_event_reactions_update_collection_invalid_artifact_filters(self):
        """A failure to filter artifacts is an error."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.get_test_user(),
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
        )
        self.create_collection(
            name=f"workflow-{wr.id}",
            category=CollectionCategory.WORKFLOW_INTERNAL,
        )
        action = ActionUpdateCollectionWithArtifacts.parse_obj(
            {
                "artifact_filters": {"xxx": "yyy"},
                "collection": "internal@collections",
            }
        )
        wr.event_reactions = EventReactions(on_success=[action])
        with self.assertLogsContains(
            "Invalid update-collection-with-artifacts artifact_filters",
            logger="debusine.db.models",
            level=logging.ERROR,
        ):
            wr.mark_completed(WorkRequest.Results.SUCCESS)

    def test_event_reactions_update_collection_invalid_variables(self):
        """A failure to expand variables is an error."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.get_test_user(),
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
        )
        self.create_collection(
            name=f"workflow-{wr.id}",
            category=CollectionCategory.WORKFLOW_INTERNAL,
        )
        self.create_artifact(category=ArtifactCategory.TEST, work_request=wr)
        action = ActionUpdateCollectionWithArtifacts.parse_obj(
            {
                "artifact_filters": {"category": ArtifactCategory.TEST},
                "collection": "internal@collections",
                "name_template": "{name}",
                "variables": {"$name": "zzz"},
            }
        )
        wr.event_reactions = EventReactions(on_success=[action])
        with self.assertLogsContains(
            "Invalid update-collection-with-artifacts variables",
            logger="debusine.db.models",
            level=logging.ERROR,
        ):
            wr.mark_completed(WorkRequest.Results.SUCCESS)

    def test_event_reactions_update_collection_different_workspace(self):
        """The updated collection must be in the same workspace."""
        wr = WorkRequest.objects.create(
            workspace=default_workspace(),
            created_by=self.get_test_user(),
            task_type=TaskTypes.WORKFLOW,
            task_name="test",
        )
        different_workspace = Workspace.objects.create_with_name(name="test")
        Collection.objects.create(
            name="test_differentworkspace",
            category=CollectionCategory.WORKFLOW_INTERNAL,
            workspace=different_workspace,
        )
        action = ActionUpdateCollectionWithArtifacts.parse_obj(
            {
                "collection": (
                    "test_differentworkspace@debusine:workflow-internal"
                ),
                "artifact_filters": {"category": ArtifactCategory.TEST},
            }
        )
        wr.event_reactions = EventReactions(on_success=[action])
        with self.assertLogsContains(
            "'test_differentworkspace@debusine:workflow-internal' does not "
            "exist or is hidden",
            logger="debusine.db.models",
            level=logging.ERROR,
        ):
            wr.mark_completed(WorkRequest.Results.SUCCESS)


class WorkflowTemplateTests(TestCase):
    """Unit tests for the WorkflowTemplate class."""

    def assert_data_must_be_dictionary(
        self, wt: WorkflowTemplate, messages: list[str] | None = None
    ):
        """Check that a WorkflowTemplate's task_data is a dictionary."""
        if messages is None:
            messages = ["task data must be a dictionary"]
        with self.assertRaises(ValidationError) as raised:
            wt.full_clean()
        self.assertEqual(
            raised.exception.message_dict,
            {"task_data": messages},
        )

    def test_clean_task_data_must_be_dict(self) -> None:
        """Ensure that task_data is a dict."""
        wt = WorkflowTemplate.objects.create(
            name="test",
            workspace=default_workspace(),
            task_name="noop",
            task_data={},
        )
        wt.full_clean()

        wt.task_data = None
        self.assert_data_must_be_dictionary(wt)

        wt.task_data = ""
        self.assert_data_must_be_dictionary(wt)

        wt.task_data = 3
        self.assert_data_must_be_dictionary(wt)

        wt.task_data = []
        self.assert_data_must_be_dictionary(wt)

        wt.task_data = object()
        self.assert_data_must_be_dictionary(
            wt,
            messages=[
                'Value must be valid JSON.',
                'task data must be a dictionary',
            ],
        )

    def test_validate_task_data(self) -> None:
        """Test orchestrator validation of task data."""
        with preserve_task_registry():

            class ValidatingWorkflow(Workflow[BaseWorkflowData]):
                """Workflow used to test validation of template data."""

                @classmethod
                def validate_template_data(cls, data: dict[str, Any]) -> None:
                    """data-controlled validation."""
                    if msg := data.get("error"):
                        raise ValueError(msg)

                def populate(self) -> None:
                    """Unused abstract method from Workflow."""
                    raise NotImplementedError()

            wt = WorkflowTemplate.objects.create(
                name="test",
                workspace=default_workspace(),
                task_name="validatingworkflow",
                task_data={},
            )
            wt.full_clean()

            wt.task_data = {"error": "example message"}
            with self.assertRaises(ValidationError) as raised:
                wt.full_clean()
            self.assertEqual(
                raised.exception.message_dict,
                {"task_data": ["example message"]},
            )
