import importlib.util
import json
from pathlib import Path
import tempfile
import unittest
from unittest import mock


ROOT = Path(__file__).resolve().parents[2]
SPEC = importlib.util.spec_from_file_location(
    "ignored_workspace_benchmark",
    ROOT / "tools/ignored_workspace_benchmark.py",
)
BENCHMARK = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(BENCHMARK)


class BenchmarkAdmissionTests(unittest.TestCase):
    def test_storage_model_accepts_null_and_uses_last_nonblank_ancestor(self) -> None:
        blockdevices = {
            "blockdevices": [
                {
                    "path": "/dev/nvme0n1",
                    "type": "disk",
                    "model": "Reference Storage",
                    "tran": "nvme",
                    "children": [
                        {
                            "path": "/dev/mapper/benchmark",
                            "type": "crypt",
                            "model": None,
                            "tran": None,
                        }
                    ],
                }
            ]
        }
        with (
            mock.patch.object(BENCHMARK, "version_output", return_value=json.dumps(blockdevices)),
            mock.patch.object(BENCHMARK.os.path, "realpath", side_effect=lambda path: path),
        ):
            information = BENCHMARK.storage_information("/dev/mapper/benchmark")

        self.assertEqual(information["storage_model"], "Reference Storage")
        self.assertEqual(information["storage_transport"], "nvme")
        self.assertEqual(information["storage_stack"], ["disk", "crypt"])

    def test_current_cgroup_swap_max_reads_root_and_nested_v2_memberships(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            temporary = Path(temporary)
            cgroup = temporary / "cgroup"
            cgroup.mkdir()
            membership = temporary / "membership"

            membership.write_text("0::/\n")
            (cgroup / "memory.swap.max").write_text("max\n")
            self.assertEqual(BENCHMARK.current_cgroup_swap_max(membership, cgroup), "max")

            nested = cgroup / "user.slice/benchmark.scope"
            nested.mkdir(parents=True)
            (nested / "memory.swap.max").write_text("0\n")
            membership.write_text("0::/user.slice/benchmark.scope\n")
            self.assertEqual(BENCHMARK.current_cgroup_swap_max(membership, cgroup), "0")

    def test_current_cgroup_swap_max_refuses_malformed_and_escaping_paths(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            temporary = Path(temporary)
            cgroup = temporary / "cgroup"
            cgroup.mkdir()
            membership = temporary / "membership"

            for value in ("2:memory:/benchmark\n", "0::/../../outside\n", "0::/ok\n1:name:/other\n"):
                with self.subTest(value=value):
                    membership.write_text(value)
                    with self.assertRaises(AssertionError):
                        BENCHMARK.current_cgroup_swap_max(membership, cgroup)

            outside = temporary / "outside"
            outside.mkdir()
            (outside / "memory.swap.max").write_text("0\n")
            (cgroup / "escape").symlink_to(outside, target_is_directory=True)
            membership.write_text("0::/escape\n")
            with self.assertRaisesRegex(AssertionError, "escapes cgroup root"):
                BENCHMARK.current_cgroup_swap_max(membership, cgroup)


if __name__ == "__main__":
    unittest.main()