import os
import unittest
from argparse import Namespace
from unittest import TestCase, mock

from trainer import TrainerArgs
from trainer.distribute import get_gpus


class TestGpusStringParsingMethods(TestCase):
    @mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"})
    def test_parse_gpus_set_in_env_var_and_args(self):
        args = Namespace(gpus="0,1")
        gpus = get_gpus(args)
        expected_value = ["0"]
        self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value))

    @mock.patch.dict(os.environ, {})
    def test_parse_gpus_set_in_args(self):
        args = Namespace(gpus="0,1")
        gpus = get_gpus(args)
        expected_value = ["0", "1"]
        self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value))

    @mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"})
    def test_parse_gpus_set_in_env_var(self):
        args = Namespace()
        gpus = get_gpus(args)
        expected_value = ["0", "1"]
        self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value))

    @mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0, 1 "})
    def test_parse_gpus_set_in_env_var_with_spaces(self):
        args = Namespace()
        gpus = get_gpus(args)
        expected_value = ["0", "1"]
        self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value))

    @mock.patch.dict(os.environ, {})
    def test_parse_gpus_set_in_args_with_spaces(self):
        args = Namespace(gpus="0, 1, 2, 3 ")
        gpus = get_gpus(args)
        expected_value = ["0", "1", "2", "3"]
        self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value))


def msg_for_test_failure(expected_value):
    return "GPU Values are expected to be " + str(expected_value)


def create_args_parser():
    parser = TrainerArgs().init_argparse(arg_prefix="")
    parser.add_argument("--gpus", type=str)
    return parser


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