# 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 the system_information."""
import subprocess
from unittest import TestCase, mock
from unittest.mock import MagicMock

import psutil

from debusine.tasks.models import WorkerType
from debusine.worker import system_information


class TestSystemInformation(TestCase):
    """Test for system_information.py functions."""

    def test_total_physical_memory(self):
        """Test total_physical_memory function."""
        total_physical_memory = system_information.total_physical_memory()

        # Assert that total_physical_memory "looks good": it is an integer
        # and at least 128 MB of memory
        self.assertIsInstance(total_physical_memory, int)
        self.assertGreaterEqual(total_physical_memory, 128 * 1024 * 1024)

        # Check that it matches the underlying implementation
        self.assertEqual(total_physical_memory, psutil.virtual_memory().total)

    def test_total_cpu_count(self):
        """Test cpu_count function."""
        cpu_count = system_information.cpu_count()

        # Assert that cpu_count "looks good": it is an integer
        # and at least 1 CPU
        self.assertIsInstance(cpu_count, int)
        self.assertGreaterEqual(cpu_count, 1)

        # Check that it matches the underlying implementation
        self.assertEqual(cpu_count, psutil.cpu_count())

    def test_architecture(self):
        """Test architecture function."""
        architecture = system_information.architecture()

        # Assert that architecture "looks good": type, length, no end of line
        self.assertIsInstance(architecture, str)
        self.assertGreaterEqual(len(architecture), 1)
        self.assertEqual(architecture, architecture.strip())

        self.assertEqual(
            architecture,
            subprocess.check_output(
                ["dpkg", "--print-architecture"], text=True
            ).strip(),
        )

    def patch_subprocess_check_output(self) -> MagicMock:
        """Patch subprocess.check_output. Return its mock."""
        patcher = mock.patch("subprocess.check_output", autospec=True)
        mocked = patcher.start()
        self.addCleanup(patcher.stop)
        return mocked

    def test_architecture_not_available(self):
        """Test architecture function: return None, dpkg not available."""
        mocked = self.patch_subprocess_check_output()
        mocked.side_effect = FileNotFoundError()
        self.assertIsNone(system_information.architecture())

        mocked.side_effect = subprocess.CalledProcessError(
            1, ["dpkg", "--print-architecture"]
        )
        self.assertIsNone(system_information.architecture())

    def test_system_metadata_without_architecture(self):
        """system_metadata() does not include architecture related keys."""
        with mock.patch(
            "debusine.worker.system_information.architecture", return_value=None
        ) as mocked:
            metadata = system_information.system_metadata(WorkerType.EXTERNAL)

        mocked.assert_called()
        self.assertNotIn("system:host_architecture", metadata)

    def test_system_metadata_external(self):
        """system_metadata() returns valid metadata for an external worker."""
        total_physical_memory = system_information.total_physical_memory()
        cpu_count = system_information.cpu_count()
        host_architecture = system_information.architecture()

        self.assertEqual(
            system_information.system_metadata(WorkerType.EXTERNAL),
            {
                "system:cpu_count": cpu_count,
                "system:total_physical_memory": total_physical_memory,
                "system:worker_type": WorkerType.EXTERNAL,
                "system:host_architecture": host_architecture,
                "system:architectures": [host_architecture],
            },
        )

    def test_system_metadata_celery(self):
        """system_metadata() returns valid metadata for a Celery worker."""
        total_physical_memory = system_information.total_physical_memory()
        cpu_count = system_information.cpu_count()
        host_architecture = system_information.architecture()

        self.assertEqual(
            system_information.system_metadata(WorkerType.CELERY),
            {
                "system:cpu_count": cpu_count,
                "system:total_physical_memory": total_physical_memory,
                "system:worker_type": WorkerType.CELERY,
                "system:host_architecture": host_architecture,
                "system:architectures": [host_architecture],
            },
        )
