# 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 worker models."""

from typing import ClassVar

from django.utils import timezone

from debusine.db.models import Token, WorkRequest, Worker
from debusine.db.models.workers import WorkerManager
from debusine.tasks.models import WorkerType
from debusine.test import TestHelpersMixin
from debusine.test.django import TestCase


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

    def test_connected(self) -> None:
        """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) -> None:
        """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.playground.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) -> None:
        """
        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.playground.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.playground.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) -> None:
        """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) -> None:
        """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) -> None:
        """
        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) -> None:
        """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) -> None:
        """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) -> None:
        """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) -> None:
        """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) -> None:
        """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) -> None:
        """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) -> None:
        """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) -> None:
        """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) -> None:
        """Test mark_connect method."""
        time_before = timezone.now()
        self.assertIsNone(self.worker.connected_at)

        self.worker.mark_connected()
        assert self.worker.connected_at is not None

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

    def test_mark_disconnected(self) -> None:
        """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) -> None:
        """Test connected method."""
        self.assertFalse(self.worker.connected())

        self.worker.connected_at = timezone.now()

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

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

        work_request = self.playground.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) -> None:
        """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.playground.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_running_work_requests(self):
        """The running_work_requests return expected work requests."""
        worker_1 = self.playground.create_worker()
        worker_1.concurrency = 3
        worker_1.save()

        worker_2 = self.playground.create_worker()

        work_request_1 = self.playground.create_work_request(
            worker=worker_1, mark_running=True
        )
        work_request_2 = self.playground.create_work_request(
            worker=worker_1, mark_running=True
        )
        self.playground.create_work_request(worker=worker_1, mark_running=False)
        self.playground.create_work_request(worker=worker_2, mark_running=True)

        self.assertQuerySetEqual(
            worker_1.running_work_requests(), [work_request_1, work_request_2]
        )

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

    def test_metadata_with_conflict(self) -> None:
        """
        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) -> None:
        """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) -> None:
        """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)
        assert self.worker.dynamic_metadata_updated_at is not None

        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) -> None:
        """Test WorkerTests.__str__."""
        self.assertEqual(
            self.worker.__str__(),
            f"Id: {self.worker.id} Name: {self.worker.name}",
        )
