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

"""Tests for the worker views."""

import datetime as dt
import logging
import secrets
from collections.abc import Generator
from typing import Any, ClassVar, override
from unittest import mock

import pytest
from django.db import OperationalError, transaction
from django.db.models import Max
from django.test import override_settings
from django.urls import reverse
from django.utils import timezone
from rest_framework import status

from debusine.artifacts.models import RuntimeStatistics
from debusine.db.models import Token, WorkRequest, Worker
from debusine.db.models.work_requests import (
    MAX_AUTOMATIC_RETRIES,
    WorkRequestQuerySet,
)
from debusine.db.tests.utils import RunInParallelTransaction, raise_db_exception
from debusine.project.celery import celery_tasks_in_progress
from debusine.server.views.rest import (
    IsWorkRequestTokenAuthenticated,
    IsWorkerTokenAuthenticated,
)
from debusine.server.views.tests.test_work_requests import WorkRequestTestCase
from debusine.server.views.workers import (
    GetNextWorkRequestView,
    UpdateWorkRequestAsCompletedView,
    UpdateWorkerDynamicMetadataView,
    update_work_request_as_completed,
)
from debusine.server.workflows.models import WorkRequestWorkflowData
from debusine.tasks.models import (
    OutputData,
    OutputDataError,
    SbuildData,
    SbuildInput,
    WorkerType,
)
from debusine.test.django import TestCase, TransactionTestCase


class RegisterViewTests(TestCase):
    """Tests for the RegisterView class."""

    def test_create_token_and_worker(self) -> None:
        """Token and Worker are created by the view."""
        self.assertQuerySetEqual(Token.objects.all(), [])
        self.assertQuerySetEqual(Worker.objects.all(), [])

        time_start = timezone.now()

        token_key = secrets.token_hex(32)

        data = {'token': token_key, 'fqdn': 'worker-bee.lan'}

        response = self.client.post(reverse('api:register'), data)

        self.assertEqual(response.status_code, status.HTTP_201_CREATED)

        token = Token.objects.get_token_or_none(token_key=token_key)
        assert token is not None

        worker = token.worker

        # Assert created worker
        self.assertEqual(worker.name, 'worker-bee-lan')
        self.assertGreaterEqual(worker.registered_at, time_start)
        self.assertLessEqual(worker.registered_at, timezone.now())
        self.assertIsNone(worker.connected_at)
        self.assertEqual(worker.worker_type, WorkerType.EXTERNAL)

        # Assert created token
        self.assertGreaterEqual(token.created_at, time_start)
        self.assertLessEqual(token.created_at, timezone.now())
        self.assertIsNone(token.user)
        self.assertEqual(token.comment, '')
        self.assertFalse(token.enabled)

    def test_create_token_and_worker_duplicate_name(self) -> None:
        """Token is created and Worker disambiguated if needed."""
        token_1 = self.playground.create_worker_token(set_worker=False)
        Worker.objects.create_with_fqdn('worker-lan', token=token_1)

        token_key = secrets.token_hex(32)
        data = {'token': token_key, 'fqdn': 'worker.lan'}
        response = self.client.post(reverse('api:register'), data)

        self.assertEqual(response.status_code, status.HTTP_201_CREATED)

        self.assertTrue(Worker.objects.filter(name='worker-lan-2').exists())

    def test_create_token_and_worker_signing(self) -> None:
        """Register a signing worker."""
        self.assertQuerySetEqual(Token.objects.all(), [])
        self.assertQuerySetEqual(Worker.objects.all(), [])

        time_start = timezone.now()

        token_key = secrets.token_hex(32)

        data = {
            'token': token_key,
            'fqdn': 'worker-bee.lan',
            'worker_type': WorkerType.SIGNING,
        }

        response = self.client.post(reverse('api:register'), data)

        self.assertEqual(response.status_code, status.HTTP_201_CREATED)

        token = Token.objects.get_token_or_none(token_key=token_key)
        assert token is not None

        worker = token.worker

        # Assert created worker
        self.assertEqual(worker.name, 'worker-bee-lan')
        self.assertGreaterEqual(worker.registered_at, time_start)
        self.assertLessEqual(worker.registered_at, timezone.now())
        self.assertIsNone(worker.connected_at)
        self.assertEqual(worker.worker_type, WorkerType.SIGNING)

        # Assert created token
        self.assertGreaterEqual(token.created_at, time_start)
        self.assertLessEqual(token.created_at, timezone.now())
        self.assertIsNone(token.user)
        self.assertEqual(token.comment, '')

    def test_invalid_data(self) -> None:
        """Request is refused if we have bad data."""
        data = {
            'token': secrets.token_hex(128),  # Too long
            'fqdn': 'worker.lan',
        }

        response = self.client.post(reverse('api:register'), data)

        self.assertResponseProblem(
            response,
            "Cannot deserialize worker",
            validation_errors_pattern=(
                r"Ensure this field has no more than .* characters"
            ),
        )
        self.assertIsInstance(response.json()["validation_errors"], dict)

        self.assertFalse(Worker.objects.filter(name='worker-lan').exists())

    def test_activation_token(self) -> None:
        """Register using an activation token."""
        activation_token = self.playground.create_worker_activation_token()
        worker = activation_token.activating_worker
        assert worker is not None
        token_key = secrets.token_hex(32)
        data = {"token": token_key, "fqdn": "worker.lan"}

        response = self.client.post(
            reverse("api:register"),
            data,
            headers={"token": activation_token.key},
        )

        self.assertEqual(response.status_code, status.HTTP_201_CREATED)
        worker.refresh_from_db()
        self.assertIsNone(worker.activation_token)
        token = Token.objects.get_token_or_none(token_key=token_key)
        assert token is not None
        self.assertEqual(token.worker, worker)
        self.assertTrue(token.enabled)
        self.assertFalse(
            Token.objects.get_token_or_none(token_key=activation_token.key)
        )

    def test_invalid_activation_token(self) -> None:
        """An invalid activation token is rejected outright."""
        token_key = secrets.token_hex(32)
        data = {"token": token_key, "fqdn": "worker.lan"}

        response = self.client.post(
            reverse("api:register"), data, headers={"token": token_key}
        )

        self.assertResponseProblem(
            response, title="Invalid worker activation token"
        )
        self.assertFalse(Token.objects.get_token_or_none(token_key=token_key))


class GetNextWorkRequestViewTests(WorkRequestTestCase):
    """Tests for GetNextWorkRequestView."""

    token: ClassVar[Token]
    worker: ClassVar[Worker]

    @override
    @classmethod
    def setUpTestData(cls) -> None:
        super().setUpTestData()
        cls.worker = cls.playground.create_worker(fqdn="worker-test")
        assert cls.worker.token is not None
        cls.token = cls.worker.token

    def create_work_request(
        self, created_at: dt.datetime, **kwargs: Any
    ) -> WorkRequest:
        """Return a new work_request as specified by the parameters."""
        environment_item = self.playground.create_debian_environment()
        assert environment_item.artifact is not None
        source_artifact = self.playground.create_source_artifact()
        task_data = SbuildData(
            input=SbuildInput(source_artifact=source_artifact.id),
            build_architecture="test",
            environment=environment_item.artifact.id,
        )

        work_request = self.playground.create_worker_task(
            worker=self.worker,
            task_name="sbuild",
            task_data=task_data,
            **kwargs,
        )

        work_request.created_at = created_at

        work_request.save()

        return work_request

    def test_check_permissions(self) -> None:
        """Only authenticated requests are processed by the view."""
        self.assertIn(
            IsWorkerTokenAuthenticated,
            GetNextWorkRequestView.permission_classes,
        )

    def test_completion_in_progress(self) -> None:
        """A work request that's pending completion is left alone."""
        work_request = self.create_work_request(
            dt.datetime.now(dt.UTC) - dt.timedelta(minutes=10),
            mark_running=True,
        )
        work_request_pending = self.create_work_request(
            dt.datetime.now(dt.UTC), mark_pending=True
        )
        self.playground.create_build_log_artifact(work_request=work_request)
        self.assertFalse(
            celery_tasks_in_progress(
                "debusine.server.views.workers.update_work_request_as_completed"
            ).exists()
        )
        token = self.playground.create_work_request_token(work_request)

        with self.captureOnCommitCallbacks(execute=True):
            response = self.client.put(
                reverse(
                    "api:work-request-completed",
                    kwargs={"work_request_id": work_request.id},
                ),
                data={"result": WorkRequest.Results.SUCCESS},
                headers={"token": token.key},
                content_type="application/json",
            )

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertTrue(
            celery_tasks_in_progress(
                "debusine.server.views.workers.update_work_request_as_completed"
            ).exists()
        )

        with self.assertNoLogs(
            logger="debusine.server.views.workers", level=logging.DEBUG
        ):
            response = self.client.post(
                reverse("api:work-request-get-next"),
                headers={"token": self.token.key},
            )

        self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)
        self.assertFalse(hasattr(work_request, "superseded"))

        # Pretend that the Celery task completed normally.
        with (
            override_settings(PGTRANSACTION_RETRY=0),
            mock.patch("pgtransaction.Atomic.execute_set_transaction_modes"),
        ):
            celery_result = update_work_request_as_completed.apply(
                args=(work_request.id,),
                kwargs={
                    "result": WorkRequest.Results.SUCCESS,
                    "output_data": None,
                },
            )
            self.assertIsNone(celery_result.get())
        celery_tasks_in_progress(
            "debusine.server.views.workers.update_work_request_as_completed"
        ).delete()

        with self.assertNoLogs(
            logger="debusine.server.views.workers", level=logging.DEBUG
        ):
            response = self.client.post(
                reverse("api:work-request-get-next"),
                headers={"token": self.token.key},
            )

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.SUCCESS)
        self.assertFalse(hasattr(work_request, "superseded"))
        self.check_response_for_work_request(
            response, work_request_pending, with_token=True
        )

    def test_get_running_work_request(self) -> None:
        """A running work request with no output is simply retried."""
        work_request = self.playground.create_worker_task(
            mark_running=True, worker=self.worker
        )

        with self.assertLogsContains(
            f"Re-attempting work request {work_request.id}",
            logger="debusine.server.views.workers",
            level=logging.DEBUG,
        ):
            response = self.client.post(
                reverse('api:work-request-get-next'),
                headers={"token": self.token.key},
            )

        # The running work request was retried
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.RUNNING)
        self.assertFalse(hasattr(work_request, "superseded"))
        self.assertEqual(work_request.workflow_data.retry_count, 1)
        self.check_response_for_work_request(
            response, work_request, with_token=True
        )

    def test_get_running_work_request_retry_limit(self) -> None:
        """A work request is not simply retried after the automatic limit."""
        work_request = self.playground.create_worker_task(
            mark_running=True,
            worker=self.worker,
            workflow_data=WorkRequestWorkflowData(
                retry_count=MAX_AUTOMATIC_RETRIES
            ),
        )

        with self.assertLogsContains(
            (
                f"Cannot retry previously-running work request "
                f"{work_request.id}: automatic retry limit exceeded."
            ),
            logger="debusine.server.views.workers",
            level=logging.DEBUG,
        ):
            response = self.client.post(
                reverse('api:work-request-get-next'),
                headers={"token": self.token.key},
            )

        # The running work request was not retried
        self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
        work_request.refresh_from_db()
        self.assertEqual(work_request.status, WorkRequest.Statuses.ABORTED)
        self.assertEqual(
            work_request.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=(
                            f"No worker managed to complete work request "
                            f"{work_request.id}. Automatic retry limit "
                            f"exceeded."
                        ),
                        code="work-request-retry-limit-exceeded",
                    )
                ]
            ),
        )
        self.assertFalse(hasattr(work_request, "superseded"))

    def test_get_running_partially_completed_work_request(self) -> None:
        """A partially-completed work request is aborted and retried."""
        # Create WorkRequest pending
        work_request_pending = self.create_work_request(
            dt.datetime(2022, 1, 5, 10, 13, 20, 204242, dt.UTC),
            mark_pending=True,
        )

        # Create WorkRequest running with output
        work_request_running = self.create_work_request(
            dt.datetime(2022, 1, 5, 11, 14, 22, 242178, dt.UTC),
            mark_running=True,
        )
        self.playground.create_build_log_artifact(
            work_request=work_request_running
        )

        # Request
        with self.assertLogs(
            logger="debusine.server.views.workers", level=logging.DEBUG
        ) as log:
            response = self.client.post(
                reverse("api:work-request-get-next"),
                headers={"token": self.token.key},
            )

        self.assertEqual(
            log.output,
            [
                f"INFO:debusine.server.views.workers:"
                f"Retrying stuck work request {work_request_running.pk}"
            ],
        )

        # The running work request was retried, though not yet assigned to a
        # worker; we got the pending one.
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        work_request_running.refresh_from_db()
        self.assertEqual(
            work_request_running.status, WorkRequest.Statuses.ABORTED
        )
        self.assertTrue(hasattr(work_request_running, "superseded"))
        self.assertIsNone(work_request_running.superseded.worker)
        self.assertEqual(work_request_running.workflow_data.retry_count, 0)
        self.assertEqual(
            work_request_running.superseded.workflow_data.retry_count, 1
        )
        self.check_response_for_work_request(
            response, work_request_pending, with_token=True
        )

    def test_get_running_preserves_parent_workflow(self) -> None:
        """The parent of a partially-completed work request is undisturbed."""
        workflow = self.playground.create_workflow(mark_running=True)
        # Allow dependency failures so that aborting all children of the
        # workflow would normally mark it as completed.
        workflow.workflow_data = WorkRequestWorkflowData(
            allow_dependency_failures=True
        )
        workflow.save()

        work_request_running = workflow.create_child_worker("noop")
        self.playground.advance_work_request(
            work_request_running, mark_running=True, worker=self.worker
        )
        self.playground.create_build_log_artifact(
            work_request=work_request_running
        )

        # Request
        with self.assertLogs(
            logger="debusine.server.views.workers", level=logging.DEBUG
        ) as log:
            response = self.client.post(
                reverse("api:work-request-get-next"),
                headers={"token": self.token.key},
            )

        self.assertEqual(
            log.output,
            [
                f"INFO:debusine.server.views.workers:"
                f"Retrying stuck work request {work_request_running.pk}"
            ],
        )

        # The running work request was retried, though not yet assigned to a
        # worker.
        self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
        work_request_running.refresh_from_db()
        self.assertEqual(
            work_request_running.status, WorkRequest.Statuses.ABORTED
        )
        self.assertTrue(hasattr(work_request_running, "superseded"))
        self.assertIsNone(work_request_running.superseded.worker)
        self.assertEqual(work_request_running.workflow_data.retry_count, 0)
        self.assertEqual(
            work_request_running.superseded.workflow_data.retry_count, 1
        )
        # The parent workflow is still running.
        workflow.refresh_from_db()
        self.assertEqual(workflow.status, WorkRequest.Statuses.RUNNING)

    def test_get_running_work_request_cannot_retry(self) -> None:
        """If a running work request cannot be retried, it is logged."""
        work_request_pending = self.create_work_request(
            dt.datetime(2022, 1, 5, 10, 13, 20, 204242, dt.UTC),
            mark_pending=True,
        )
        work_request_running = self.create_work_request(
            dt.datetime(2022, 1, 5, 11, 14, 22, 242178, dt.UTC),
            mark_running=True,
        )
        work_request_running.task_data = {"invalid": True}
        work_request_running.configured_task_data = {"invalid": True}
        work_request_running.save()
        self.playground.create_build_log_artifact(
            work_request=work_request_running
        )

        with self.assertLogs(
            logger="debusine.server.views.workers", level=logging.DEBUG
        ) as log:
            response = self.client.post(
                reverse('api:work-request-get-next'),
                headers={"token": self.token.key},
            )
            self.assertAPIResponseOk(response)

        self.assertEqual(
            log.output,
            [
                f"INFO:debusine.server.views.workers:"
                f"Retrying stuck work request {work_request_running.pk}",
                "DEBUG:debusine.server.views.workers:"
                "Cannot retry previously-running work request: "
                "Task dependencies cannot be satisfied",
            ],
        )

        # The running work request was aborted but could not be retried; we
        # got the pending one.
        work_request_running.refresh_from_db()
        self.assertEqual(
            work_request_running.status, WorkRequest.Statuses.ABORTED
        )
        self.assertFalse(hasattr(work_request_running, "superseded"))
        self.check_response_for_work_request(
            response, work_request_pending, with_token=True
        )

    def test_get_work_request_change_status(self) -> None:
        """Get a WorkRequest change the status to running."""
        work_request = self.create_work_request(
            dt.datetime(2022, 1, 5, 10, 13, 20, 204242, dt.UTC),
            mark_pending=True,
        )

        # Request
        response = self.client.post(
            reverse('api:work-request-get-next'),
            headers={"token": self.token.key},
        )

        # Assert that we got the WorkRequest running
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.check_response_for_work_request(
            response, work_request, with_token=True
        )

        work_request.refresh_from_db()

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

    def test_get_work_request_configured_task_data(self) -> None:
        """Get a WorkRequest shows configured task data."""
        work_request = self.create_work_request(
            dt.datetime(2022, 1, 5, 10, 13, 20, 204242, dt.UTC),
            mark_pending=True,
        )
        work_request.task_data = {"test": 1}
        work_request.configured_task_data = {"test": 2}
        work_request.save()

        # Request
        response = self.client.post(
            reverse('api:work-request-get-next'),
            headers={"token": self.token.key},
        )

        # Assert that we got the WorkRequest running
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.check_response_for_work_request(
            response, work_request, with_token=True
        )
        self.assertEqual(response.json()["task_data"], {"test": 2})

    def test_get_older_pending(self) -> None:
        """View return the older pending work request."""
        # Create WorkRequest pending
        work_request_pending_01 = self.create_work_request(
            dt.datetime(2022, 1, 5, 10, 13, 20, 204242, dt.UTC),
            mark_pending=True,
        )

        work_request_pending_02 = self.create_work_request(
            dt.datetime(2022, 1, 5, 10, 13, 22, 204242, dt.UTC),
            mark_pending=True,
        )
        # Request
        response = self.client.post(
            reverse('api:work-request-get-next'),
            headers={"token": self.token.key},
        )

        self.check_response_for_work_request(
            response, work_request_pending_01, with_token=True
        )

        # work_request_pending_01 status changed. Let's move it back to
        # pending. This is not an allowed transition but helps here to test
        # that clients will receive the older pending
        work_request_pending_01.status = WorkRequest.Statuses.PENDING

        # Change created_at to force a change in the order that debusine
        # sends the WorkRequests to the client
        work_request_pending_01.created_at = dt.datetime(
            2022, 1, 5, 10, 13, 24, 204242, dt.UTC
        )

        work_request_pending_01.save()

        # New request to check the new order
        response = self.client.post(
            reverse('api:work-request-get-next'),
            headers={"token": self.token.key},
        )

        self.check_response_for_work_request(
            response, work_request_pending_02, with_token=True
        )

    def test_get_no_work_request(self) -> None:
        """Assert debusine sends HTTP 204 when nothing assigned to Worker."""
        response = self.client.post(
            reverse('api:work-request-get-next'),
            headers={"token": self.token.key},
        )

        self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)

    def test_authentication_credentials_not_provided(self) -> None:
        """Assert Token is required to use the endpoint."""
        # This is to double-check that IsWorkerTokenAuthenticated (inclusion
        # is verified in test_check_permissions) is returning what it should do.
        # All the views with IsWorkerTokenAuthenticated behave the same way
        # IsWorkerTokenAuthenticated is tested in
        # IsWorkerTokenAuthenticatedTests
        response = self.client.post(reverse('api:work-request-get-next'))
        self.assertEqual(
            response.json(),
            {
                "title": "Error",
                "detail": "Authentication credentials were not provided.",
            },
        )
        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)


class GetNextWorkRequestViewTransactionTests(TransactionTestCase):
    """Test transactional behaviour of GetNextWorkRequestView."""

    def test_locks_workflow(self) -> None:
        """The view locks all work requests in the affected workflows."""
        with transaction.atomic():
            workflows = [self.playground.create_workflow() for _ in range(5)]
            work_requests = [
                workflow.create_child_worker("noop")
                for workflow in (
                    workflows[0],
                    workflows[0],
                    workflows[1],
                    workflows[1],
                    workflows[2],
                    workflows[3],
                )
            ]
            # workflows[0] and workflows[1] have running work requests.
            for work_request in (work_requests[0], work_requests[2]):
                self.playground.advance_work_request(
                    work_request, mark_running=True
                )
                self.playground.create_build_log_artifact(
                    work_request=work_request
                )
            assert work_requests[0].worker is not None
            assert work_requests[0].worker.token is not None
            # workflows[2] and workflows[3] have pending work requests
            # assigned to the same worker as work_requests[0], with the one
            # in workflows[2] having higher priority.
            work_requests[4].priority_adjustment += 1
            work_requests[4].save()
            for work_request in (work_requests[4], work_requests[5]):
                self.playground.advance_work_request(
                    work_request,
                    mark_pending=True,
                    worker=work_requests[0].worker,
                )

        with transaction.atomic():
            response = self.client.post(
                reverse("api:work-request-get-next"),
                headers={"token": work_requests[0].worker.token.key},
            )

            self.assertEqual(response.status_code, status.HTTP_200_OK)
            # Each work request in each of the affected workflows is locked,
            # but not unaffected work requests.
            self.assert_work_requests_locked(
                locked=[
                    workflows[0],
                    work_requests[0],
                    work_requests[1],
                    workflows[2],
                    work_requests[4],
                ],
                unlocked=[
                    workflows[1],
                    work_requests[2],
                    work_requests[3],
                    workflows[3],
                    work_requests[5],
                    workflows[4],
                ],
            )

    def test_get_work_request_locked(self) -> None:
        """If the workflow is already locked, return HTTP 423."""
        with transaction.atomic():
            workflow = self.playground.create_workflow(mark_running=True)
            work_requests = [
                workflow.create_child_worker("noop") for _ in range(2)
            ]
            self.playground.advance_work_request(
                work_requests[0], mark_pending=True, assign_new_worker=True
            )
            assert work_requests[0].worker is not None
            assert work_requests[0].worker.token is not None

        thread = RunInParallelTransaction(
            lambda: WorkRequest.objects.select_for_update().get(
                id=work_requests[1].id
            )
        )
        thread.start_transaction()

        try:
            with transaction.atomic():
                with self.assertLogsContains(
                    f"Failed to lock work requests assigned to "
                    f"{work_requests[0].worker.name}",
                    logger="debusine.server.views.workers",
                    level=logging.DEBUG,
                ):
                    response = self.client.post(
                        reverse("api:work-request-get-next"),
                        headers={"token": work_requests[0].worker.token.key},
                    )

                self.assertResponseProblem(
                    response,
                    title="Retry",
                    detail_pattern=(
                        "Could not acquire workflow locks; retry later"
                    ),
                    status_code=status.HTTP_423_LOCKED,
                )
                work_requests[0].refresh_from_db()
                self.assertEqual(
                    work_requests[0].status, WorkRequest.Statuses.PENDING
                )
        finally:
            thread.stop_transaction()


class UpdateWorkRequestAsCompletedViewTests(TransactionTestCase):
    """Tests for UpdateWorkRequestAsCompletedView class and its Celery task."""

    token: Token
    worker: Worker
    work_request: WorkRequest

    @override
    def setUp(self) -> None:
        super().setUp()
        self.worker = self.playground.create_worker(fqdn="worker-test")

        environment_item = self.playground.create_debian_environment()
        assert environment_item.artifact is not None
        source_artifact = self.playground.create_source_artifact()
        task_data = SbuildData(
            input=SbuildInput(source_artifact=source_artifact.id),
            build_architecture="test",
            environment=environment_item.artifact.id,
        )

        self.work_request = self.playground.create_worker_task(
            worker=self.worker,
            task_name='sbuild',
            task_data=task_data,
            mark_running=True,
        )
        self.token = self.playground.create_work_request_token(
            self.work_request
        )

    def test_check_permissions(self) -> None:
        """Only authenticated requests are processed by the view."""
        self.assertIn(
            IsWorkRequestTokenAuthenticated,
            UpdateWorkRequestAsCompletedView.permission_classes,
        )

    def put_api_work_request_completed(
        self, result: WorkRequest.Results, output_data: OutputData | None = None
    ) -> None:
        """Assert Worker can update WorkRequest result to completed-success."""
        self.playground.advance_work_request(
            self.work_request, mark_running=True
        )

        data: dict[str, Any] = {"result": result}
        if output_data is not None:
            data["output_data"] = output_data.model_dump(
                mode="json", exclude_unset=True
            )
        with mock.patch(
            "debusine.server.views.workers.update_work_request_as_completed"
            ".apply_async"
        ) as mock_apply_async:
            response = self.client.put(
                reverse(
                    "api:work-request-completed",
                    kwargs={"work_request_id": self.work_request.id},
                ),
                data=data,
                headers={"token": self.token.key},
                content_type="application/json",
            )

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        mock_apply_async.assert_called_once_with(
            args=(self.work_request.id,),
            kwargs={"result": result, "output_data": data.get("output_data")},
            task_id=mock.ANY,
        )

        # The work request status is not changed yet; that work is deferred
        # to a Celery task.
        self.work_request.refresh_from_db()
        self.assertEqual(self.work_request.status, WorkRequest.Statuses.RUNNING)

        celery_result = update_work_request_as_completed.apply(
            args=(self.work_request.id,),
            kwargs={
                "result": data["result"],
                "output_data": data.get("output_data"),
            },
        )

        self.assertIsNone(celery_result.get())
        self.work_request.refresh_from_db()
        self.assertEqual(
            self.work_request.status, WorkRequest.Statuses.COMPLETED
        )
        self.assertEqual(self.work_request.result, result)
        self.assertEqual(self.work_request.output_data, output_data)

    def test_update_completed_result_is_success(self) -> None:
        """Assert Worker can update WorkRequest result to completed-success."""
        self.put_api_work_request_completed(WorkRequest.Results.SUCCESS)

    def test_update_completed_result_is_success_with_output_data(self) -> None:
        """Assert Worker can update WorkRequest result with output data."""
        self.put_api_work_request_completed(
            WorkRequest.Results.SUCCESS,
            OutputData(runtime_statistics=RuntimeStatistics(duration=60)),
        )

    def test_update_completed_result_is_failure(self) -> None:
        """Assert Worker can update WorkRequest result to completed-error."""
        self.put_api_work_request_completed(WorkRequest.Results.ERROR)

    def test_update_work_request_id_not_found(self) -> None:
        """Assert API returns HTTP 404 for a non-existing WorkRequest id."""
        max_id = WorkRequest.objects.aggregate(Max('id'))['id__max']

        response = self.client.put(
            reverse(
                'api:work-request-completed',
                kwargs={'work_request_id': max_id + 1},
            ),
            data={"result": WorkRequest.Results.SUCCESS},
            headers={"token": self.token.key},
            content_type="application/json",
        )

        self.assertResponseProblem(
            response,
            "Work request not found",
            status_code=status.HTTP_404_NOT_FOUND,
        )

    def test_celery_work_request_id_not_found(self) -> None:
        """The Celery task fails if the work request does not exist."""
        max_id = WorkRequest.objects.aggregate(Max("id"))["id__max"]

        with self.assertLogsContains(
            f"Work request {max_id + 1} does not exist",
            logger="debusine.server.views.workers",
            level=logging.ERROR,
        ):
            result = update_work_request_as_completed.apply(
                args=(max_id + 1,),
                kwargs={"result": WorkRequest.Results.SUCCESS},
            )

        with self.assertRaises(WorkRequest.DoesNotExist):
            result.get()

    def test_celery_work_request_already_completed(self) -> None:
        """The Celery task fails if the work request was already completed."""
        self.work_request.mark_completed(WorkRequest.Results.SUCCESS)

        with self.assertLogsContains(
            f"Failed to mark work request {self.work_request.id} as completed",
            logger="debusine.server.views.workers",
            level=logging.ERROR,
        ):
            result = update_work_request_as_completed.apply(
                args=(self.work_request.id,),
                kwargs={"result": WorkRequest.Results.SUCCESS},
            )

        self.assertIsNone(result.get())
        self.work_request.refresh_from_db()
        self.assertFalse(hasattr(self.work_request, "superseded"))
        self.assertEqual(
            self.work_request.status, WorkRequest.Statuses.COMPLETED
        )
        self.assertEqual(self.work_request.result, WorkRequest.Results.SUCCESS)

    def test_celery_work_request_failure_when_running(self) -> None:
        """The Celery task prevents the worker getting stuck on failure."""
        self.playground.advance_work_request(
            self.work_request, mark_running=True
        )

        with (
            self.assertLogsContains(
                f"Failed to mark work request {self.work_request.id} as"
                f" completed",
                logger="debusine.server.views.workers",
                level=logging.ERROR,
            ),
            mock.patch.object(
                WorkRequest,
                "mark_completed",
                side_effect=lambda *args, **kwargs: raise_db_exception(
                    "serialization_failure"
                ),
            ),
        ):
            result = update_work_request_as_completed.apply(
                args=(self.work_request.id,),
                kwargs={"result": WorkRequest.Results.SUCCESS},
            )

        self.assertIsNone(result.get())
        self.work_request.refresh_from_db()
        assert self.work_request.superseded is not None
        self.assertEqual(
            self.work_request.superseded.workflow_data.retry_count, 1
        )

    def test_update_failure_invalid_result(self) -> None:
        """Assert Worker get HTTP 400 and error for invalid result field."""
        self.playground.advance_work_request(
            self.work_request, mark_running=True
        )

        response = self.client.put(
            reverse(
                'api:work-request-completed',
                kwargs={'work_request_id': self.work_request.id},
            ),
            data={"result": "something not recognised"},
            headers={"token": self.token.key},
            content_type="application/json",
        )

        self.assertResponseProblem(
            response,
            "Cannot change work request as completed",
            validation_errors_pattern="something not recognised",
        )

    def test_unauthorized(self) -> None:
        """Assert worker cannot modify a task of another worker."""
        another_worker = self.playground.create_worker(
            fqdn="another-worker-test"
        )

        environment_item = self.playground.create_debian_environment()
        assert environment_item.artifact is not None
        source_artifact = self.playground.create_source_artifact()
        task_data = SbuildData(
            input=SbuildInput(source_artifact=source_artifact.id),
            build_architecture="test",
            environment=environment_item.artifact.id,
        )

        another_work_request = self.playground.create_worker_task(
            worker=another_worker,
            task_name='sbuild',
            task_data=task_data,
            mark_pending=True,
        )

        response = self.client.put(
            reverse(
                'api:work-request-completed',
                kwargs={'work_request_id': another_work_request.id},
            ),
            headers={"token": self.token.key},
            data={"result": WorkRequest.Results.SUCCESS},
            content_type="application/json",
        )

        self.assertResponseProblem(
            response,
            "Work request not found",
            status_code=status.HTTP_404_NOT_FOUND,
        )

    def test_locks_workflow(self) -> None:
        """The Celery task locks all work requests in the workflow."""
        with transaction.atomic():
            workflow = self.playground.create_workflow()
            work_requests = [
                workflow.create_child_worker("noop") for _ in range(2)
            ]
            self.playground.advance_work_request(
                work_requests[0], mark_running=True
            )
            assert work_requests[0].worker is not None
            assert work_requests[0].worker.token is not None
            unrelated = self.playground.create_worker_task()

        with (
            override_settings(PGTRANSACTION_RETRY=0),
            mock.patch("pgtransaction.Atomic.execute_set_transaction_modes"),
            transaction.atomic(),
        ):
            result = update_work_request_as_completed.apply(
                args=(work_requests[0].id,),
                kwargs={"result": WorkRequest.Results.SUCCESS},
            )
            self.assertIsNone(result.get())

            work_requests[0].refresh_from_db()
            self.assertEqual(
                work_requests[0].status, WorkRequest.Statuses.COMPLETED
            )
            self.assertEqual(
                work_requests[0].result, WorkRequest.Results.SUCCESS
            )
            self.assert_work_requests_locked(
                locked=[workflow, *work_requests], unlocked=[unrelated]
            )

    def test_deadlocks(self) -> None:
        """The Celery task retries automatically if it deadlocks."""
        with transaction.atomic():
            workflow = self.playground.create_workflow()
            work_requests = [
                workflow.create_child_worker("noop") for _ in range(2)
            ]
            for work_request in work_requests:
                self.playground.advance_work_request(
                    work_request, mark_running=True
                )

        def lock_workflows_side_effect() -> Generator[Any]:
            with pytest.raises(OperationalError) as exc_info:
                raise_db_exception("deadlock_detected")
            yield exc_info.value

            while True:
                yield mock.DEFAULT

        with self.spy_object(
            WorkRequestQuerySet,
            "lock_workflows_for_update",
            side_effect=lock_workflows_side_effect(),
        ) as mock_lock_workflows_for_update:
            result = update_work_request_as_completed.apply(
                args=(work_requests[0].id,),
                kwargs={"result": WorkRequest.Results.SUCCESS},
            )
            self.assertIsNone(result.get())

        self.assertEqual(mock_lock_workflows_for_update.call_count, 2)
        work_requests[0].refresh_from_db()
        self.assertEqual(
            work_requests[0].status, WorkRequest.Statuses.COMPLETED
        )
        self.assertEqual(work_requests[0].result, WorkRequest.Results.SUCCESS)


class UpdateWorkerDynamicMetadataViewTests(TestCase):
    """Tests for UpdateWorkerDynamicMetadataView class."""

    @override
    def setUp(self) -> None:
        super().setUp()
        self.worker_01 = Worker.objects.create_with_fqdn(
            "worker-01-lan",
            token=self.playground.create_worker_token(set_worker=False),
        )
        self.worker_02 = Worker.objects.create_with_fqdn(
            "worker-02-lan",
            token=self.playground.create_worker_token(set_worker=False),
        )

    def test_check_permissions(self) -> None:
        """Only authenticated requests are processed by the view."""
        self.assertIn(
            IsWorkerTokenAuthenticated,
            UpdateWorkerDynamicMetadataView.permission_classes,
        )

    def test_update_metadata_success(self) -> None:
        """Worker's dynamic_metadata is updated."""
        assert self.worker_01.token is not None
        self.assertEqual(self.worker_01.dynamic_metadata, {})
        self.assertEqual(self.worker_02.dynamic_metadata, {})

        metadata = {"cpu_cores": 4, "ram": 16}
        response = self.client.put(
            reverse('api:worker-dynamic-metadata'),
            data=metadata,
            headers={"token": self.worker_01.token.key},
            content_type="application/json",
        )

        self.worker_01.refresh_from_db()

        self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)

        self.assertEqual(self.worker_01.dynamic_metadata, metadata)
        assert self.worker_01.dynamic_metadata_updated_at is not None
        self.assertLessEqual(
            self.worker_01.dynamic_metadata_updated_at, timezone.now()
        )

        self.assertEqual(self.worker_02.dynamic_metadata, {})
        self.assertIsNone(self.worker_02.dynamic_metadata_updated_at)
