from pathlib import Path
from tempfile import TemporaryDirectory
import unittest

from mm.cli import load_machines, parse_nvidia_smi_csv


class ConfigTests(unittest.TestCase):
    def test_loads_string_and_object_machine_entries(self) -> None:
        with TemporaryDirectory() as tmpdir:
            config_path = Path(tmpdir) / "list.yaml"
            config_path.write_text(
                """
machines:
  - dash0
  - alias: dash1
    label: training-1
""",
                encoding="utf-8",
            )

            machines = load_machines(config_path)

        self.assertEqual([machine.alias for machine in machines], ["dash0", "dash1"])
        self.assertEqual([machine.label for machine in machines], ["dash0", "training-1"])

    def test_loads_mapping_machine_entries(self) -> None:
        with TemporaryDirectory() as tmpdir:
            config_path = Path(tmpdir) / "list.yaml"
            config_path.write_text(
                """
machines:
  dash0:
  dash1:
    label: training-1
""",
                encoding="utf-8",
            )

            machines = load_machines(config_path)

        self.assertEqual([machine.alias for machine in machines], ["dash0", "dash1"])
        self.assertEqual([machine.label for machine in machines], ["dash0", "training-1"])


class NvidiaSmiParsingTests(unittest.TestCase):
    def test_parse_csv_output(self) -> None:
        output = (
            "0, NVIDIA A100-SXM4-80GB, 12000, 81920, 84, 61, 294.3, 400.0\n"
            "1, NVIDIA A100-SXM4-80GB, 0, 81920, 0, 32, 49.8, 400.0\n"
        )

        gpus = parse_nvidia_smi_csv(output)

        self.assertEqual(len(gpus), 2)
        self.assertEqual(gpus[0].index, "0")
        self.assertEqual(gpus[0].memory_used_mib, 12000)
        self.assertEqual(gpus[0].utilization_pct, 84)
        self.assertEqual(gpus[1].power_draw_w, 49.8)


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