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

"""Tests for scheduler."""

import inspect
from datetime import timedelta
from typing import Optional

from channels.db import database_sync_to_async

from django.test import TestCase, TransactionTestCase, override_settings
from django.utils import timezone

from debusine.db.models import Token, WorkRequest, Worker
from debusine.db.tests.utils import RunInParallelTransaction
from debusine.server.scheduler import (
    _work_request_changed,
    _worker_changed,
    schedule,
    schedule_for_worker,
)
from debusine.tasks.models import WorkerType
from debusine.tasks.sbuild import Sbuild
from debusine.test import TestHelpersMixin


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

    SAMPLE_TASK_DATA = {
        "input": {
            "source_artifact_id": 10,
        },
        "distribution": "unstable",
        "backend": "schroot",
        "host_architecture": "amd64",
        "build_components": [
            "any",
            "all",
        ],
        "sbuild_options": [
            "--post-build-commands=/usr/local/bin/post-process %SBUILD_CHANGES",
        ],
    }

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

        sbuild_chroots = "%s-%s" % (
            cls.SAMPLE_TASK_DATA['distribution'],
            cls.SAMPLE_TASK_DATA['host_architecture'],
        )

        worker.dynamic_metadata = {
            "system:worker_type": worker_type,
            "sbuild:version": Sbuild.TASK_VERSION,
            "sbuild:available": True,
            "sbuild:chroots": [sbuild_chroots],
        }

        worker.mark_connected()

        return worker

    @classmethod
    def create_work_request(cls, assign_to: Optional[Worker] = None):
        """Create and return a WorkRequest."""
        work_request = TestHelpersMixin.create_work_request(task_name='sbuild')
        work_request.task_data = cls.SAMPLE_TASK_DATA
        work_request.save()

        if assign_to:
            work_request.assign_worker(assign_to)

        return work_request

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

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


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

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

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

        self.assertEqual(schedule(), [])

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

        result = schedule()

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

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

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

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

        self.assertEqual(schedule(), [])

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

        self.assertEqual(schedule(), [])


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

    def setUp(self):
        """Initialize test case."""
        self.worker = self.create_worker(enabled=True)

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

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

        result = schedule_for_worker(self.worker)

        self.assertIsNone(result)

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

        result = schedule_for_worker(self.worker)

        self.assertIsNone(result)

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

        result = schedule_for_worker(self.worker)

        self.assertIsNone(result)

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

        result = schedule_for_worker(self.worker)

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

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

        result = schedule_for_worker(self.worker)
        work_request.refresh_from_db()

        self.assertIsNone(result)
        self.assertEqual(work_request.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(work_request.result, WorkRequest.Results.ERROR)

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

        self.assertIsNone(schedule_for_worker(self.worker))

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

        self.assertIsNone(schedule_for_worker(self.worker))

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

        result = schedule_for_worker(self.worker)

        self.assertEqual(work_request, result)

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

        result = schedule_for_worker(self.worker)

        self.assertEqual(list_of_work_requests[1], result)

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

        work_request.status = status
        work_request.save()

        self.assertIsNone(schedule_for_worker(self.worker))

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

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

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

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

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

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


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

    def setUp(self):
        """Initialize test case."""
        self.worker = self.create_worker(enabled=True)

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

        self.assertIsNone(work_request_1.worker)

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

        self.assertEqual(work_request_1.worker, work_request_1_worker)

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

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

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

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

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

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

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

        work_request_1 = await database_sync_to_async(
            self.create_work_request
        )()

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

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

        await database_sync_to_async(work_request_1.refresh_from_db)()

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

        self.assertEqual(work_request_1_worker, self.worker)

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

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

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

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

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

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


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

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

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