import unittest
import argparse

class TestPatch(unittest.TestCase):

    def setUp(self):
        self.parser = p = argparse.ArgumentParser()
        p.add_argument('positional', nargs='*')
        p.add_argument('-o', '--optional')

    def add_argument(self, lo, hi, name='--foo'):
        return self.parser.add_argument(name, nargs=(lo, hi))

    ### tests ###

    def test_add_argument1(self):
        "len(nargs) = 0"
        with self.assertRaises(ValueError):
            self.parser.add_argument('-a', nargs=() )

    def test_add_argument2(self):
        "len(nargs) = 1"
        with self.assertRaises(ValueError):
            self.parser.add_argument('-a', nargs=(1,) )

    def test_add_argument3(self):
        "len(nargs) = 2, first argument is not a valid num"
        with self.assertRaises(ValueError):
            self.add_argument('xxx', None)

    def test_add_argument4(self):
        "len(nargs) = 2, second argument is not valid num"
        with self.assertRaises(ValueError):
            self.add_argument(None, 'xxx')

    def test_add_argument5(self):
        "nargs = (lo, hi), where lo < hi"
        arg = (3, 5)
        res = arg

        opt = self.add_argument(*arg)
        self.assertEqual(opt.nargs, res)

    def test_add_argument6(self):
        "nargs = (lo, hi), where lo == hi, optimized to single number"
        arg = (12, 12)
        res = 12

        opt = self.add_argument(*arg)
        self.assertEqual(opt.nargs, res)

    def test_add_argument7(self):
        "nargs = (lo, hi), where lo > hi (error)"

        with self.assertRaises(ValueError):
            opt = self.add_argument(8, 3)

    def test_add_argument8(self):
        "nargs = (None, None) => (0, None) -- optimized to '*'"
        opt = self.add_argument(0, None)
        self.assertEqual(opt.nargs, argparse.ZERO_OR_MORE)

    def test_add_argument9(self):
        "nargs = (1, None) => optimized to '+'"
        opt = self.add_argument(1, None)
        self.assertEqual(opt.nargs, argparse.ONE_OR_MORE)

    def test_add_argument10(self):
        "nargs = (0, 1) => optimized to '?'"
        opt = self.add_argument(1, None)
        self.assertEqual(opt.nargs, argparse.ONE_OR_MORE)

    def test_add_argument11(self):
        "nargs = (None, number) => (0, number)"
        arg = (None, 87)
        res = (0, 87)

        opt = self.add_argument(*arg)
        self.assertEqual(opt.nargs, res)

    def test_parser_args_a1(self):
        "nargs=(3, 7) - consume all possible options"
        args = "aa bb cc dd ee ff gg hh".split()

        self.add_argument(3, 7)
        n = self.parser.parse_args(["--foo"] + args)

        self.assertEqual(n.foo, args[:7])
        self.assertEqual(n.positional, args[7:])

    def test_parser_args_a2(self):
        "nargs=(3, 7) - consume 4 options"
        args = "aa bb cc dd -o ee ff gg hh".split()

        self.add_argument(3, 7)
        n = self.parser.parse_args(["--foo"] + args)

        self.assertEqual(n.foo, args[:4])
        self.assertEqual(n.positional, args[6:])

    def test_parser_args_a3(self):
        "nargs=(3, 7) - can't consume at least 3 options"
        args = "aa bb -o cc dd ee ff gg hh".split()

        self.add_argument(3, 7)
        with self.assertRaises(SystemExit):
            self.parser.parse_args(["--foo"] + args)

    def test_parser_args_a4(self):
        "nargs=(3, 7) - can't consume at least 3 options (another case)"
        args = "-o aa bb cc dd ee ff gg hh".split()

        self.add_argument(3, 5)
        with self.assertRaises(SystemExit):
            self.parser.parse_args(["--foo"] + args)

    def test_parser_args_a5(self):
        "nargs=(3, None) - consume as many as possible "
        args = "aa bb cc dd ee ff gg hh ii -o jj kk".split()

        self.add_argument(3, None)
        n = self.parser.parse_args(["--foo"] + args)
        self.assertEqual(n.foo, args[:9])

    def test_parser_args_a6(self):
        "nargs=(3, None) - consume as many as possible (error)"
        args = "aa bb -o cc dd ee ff gg hh ii -o jj kk".split()

        self.add_argument(3, None)
        with self.assertRaises(SystemExit):
            self.parser.parse_args(["--foo"] + args)


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