"""
Tests for the real argument tables of each supported launcher (srun, aprun, jsrun).

These tests verify that realistic command lines split correctly into
(launcher, program) pairs using the actual ARGUMENTS dictionary defined in
each launcher module - not a synthetic test module.  A wrong nargs entry in
a table would silently pass a program argument on the launcher side, so this
catches regressions that are invisible to the generic parser tests.
"""

from shlex import split
import tests
from e4s_cl.cf.launchers import interpret


class SrunParserTest(tests.TestCase):
    """Realistic srun command lines parsed with the real slurm ARGUMENTS table."""

    def test_basic_task_count(self):
        """srun -n 4 ./app"""
        launcher, command = interpret(split('srun -n 4 ./app'))
        self.assertEqual(launcher, ['srun', '-n', '4'])
        self.assertEqual(command, ['./app'])

    def test_ntasks_long_equals(self):
        """srun --ntasks=4 ./app"""
        launcher, command = interpret(split('srun --ntasks=4 ./app'))
        self.assertIn('--ntasks=4', launcher)
        self.assertEqual(command, ['./app'])

    def test_ntasks_long_space(self):
        """srun --ntasks 4 ./app"""
        launcher, command = interpret(split('srun --ntasks 4 ./app'))
        self.assertIn('--ntasks', launcher)
        self.assertIn('4', launcher)
        self.assertEqual(command, ['./app'])

    def test_nodes_and_tasks(self):
        """srun -N 2 -n 4 ./app"""
        launcher, command = interpret(split('srun -N 2 -n 4 ./app'))
        self.assertEqual(launcher, ['srun', '-N', '2', '-n', '4'])
        self.assertEqual(command, ['./app'])

    def test_account_flag(self):
        """srun -A myaccount ./app"""
        launcher, command = interpret(split('srun -A myaccount ./app'))
        self.assertEqual(launcher, ['srun', '-A', 'myaccount'])
        self.assertEqual(command, ['./app'])

    def test_partition_flag(self):
        """srun -p gpu ./app"""
        launcher, command = interpret(split('srun -p gpu ./app'))
        self.assertIn('gpu', launcher)
        self.assertEqual(command, ['./app'])

    def test_double_dash_separator(self):
        """srun -n 4 -- ./app --bar  (program args don't bleed into launcher)

        interpret() strips '--' from both sides; everything before it is the
        launcher and everything after is the program.
        """
        launcher, command = interpret(
            ['srun', '-n', '4', '--', './app', '--bar'])
        self.assertIn('srun', launcher)
        self.assertNotIn('./app', launcher)
        self.assertEqual(command, ['./app', '--bar'])

    def test_concatenated_n(self):
        """srun -n4 ./app  (concatenated known short-arg)"""
        launcher, command = interpret(split('srun -n4 ./app'))
        self.assertIn('srun', launcher)
        # The program must not be swallowed into the launcher
        self.assertEqual(command, ['./app'])

    def test_gpu_flags(self):
        """srun --gpus-per-task=1 -n 4 ./app"""
        launcher, command = interpret(split('srun --gpus-per-task=1 -n 4 ./app'))
        self.assertIn('--gpus-per-task=1', launcher)
        self.assertEqual(command, ['./app'])

    def test_job_name(self):
        """srun --job-name=myjob -n 2 ./app"""
        launcher, command = interpret(split('srun --job-name=myjob -n 2 ./app'))
        self.assertIn('--job-name=myjob', launcher)
        self.assertEqual(command, ['./app'])

    def test_program_with_args(self):
        """srun -n 2 ./app arg1 arg2"""
        launcher, command = interpret(split('srun -n 2 ./app arg1 arg2'))
        self.assertEqual(launcher, ['srun', '-n', '2'])
        self.assertEqual(command, ['./app', 'arg1', 'arg2'])


class AprunParserTest(tests.TestCase):
    """Realistic aprun (Cray ALPS) command lines."""

    def test_basic_width(self):
        """aprun -n 4 ./app"""
        launcher, command = interpret(split('aprun -n 4 ./app'))
        self.assertEqual(launcher, ['aprun', '-n', '4'])
        self.assertEqual(command, ['./app'])

    def test_pes_per_node(self):
        """aprun -n 4 -N 2 ./app"""
        launcher, command = interpret(split('aprun -n 4 -N 2 ./app'))
        self.assertEqual(launcher, ['aprun', '-n', '4', '-N', '2'])
        self.assertEqual(command, ['./app'])

    def test_depth_flag(self):
        """aprun -n 4 -d 4 ./app"""
        launcher, command = interpret(split('aprun -n 4 -d 4 ./app'))
        self.assertEqual(launcher, ['aprun', '-n', '4', '-d', '4'])
        self.assertEqual(command, ['./app'])

    def test_double_dash(self):
        """aprun double-dash separator - interpret() strips '--' from both sides."""
        launcher, command = interpret(['aprun', '-n', '4', '--', './app'])
        self.assertIn('aprun', launcher)
        self.assertNotIn('./app', launcher)
        self.assertEqual(command, ['./app'])

    def test_memory_per_pe(self):
        """aprun -n 4 -m 512m ./app"""
        launcher, command = interpret(split('aprun -n 4 -m 512m ./app'))
        self.assertIn('-m', launcher)
        self.assertEqual(command, ['./app'])

    def test_program_with_args(self):
        """aprun -n 4 ./app --config foo"""
        launcher, command = interpret(split('aprun -n 4 ./app --config foo'))
        self.assertEqual(launcher, ['aprun', '-n', '4'])
        self.assertEqual(command, ['./app', '--config', 'foo'])


class JsrunParserTest(tests.TestCase):
    """Realistic jsrun (IBM LSF) command lines."""

    def test_basic_nrs(self):
        """jsrun -n 4 ./app"""
        launcher, command = interpret(split('jsrun -n 4 ./app'))
        self.assertEqual(launcher, ['jsrun', '-n', '4'])
        self.assertEqual(command, ['./app'])

    def test_resource_set(self):
        """jsrun --nrs 4 --cpu_per_rs 2 ./app"""
        launcher, command = interpret(
            split('jsrun --nrs 4 --cpu_per_rs 2 ./app'))
        self.assertIn('--nrs', launcher)
        self.assertIn('4', launcher)
        self.assertIn('--cpu_per_rs', launcher)
        self.assertEqual(command, ['./app'])

    def test_gpu_per_rs(self):
        """jsrun --nrs 2 --gpu_per_rs 1 ./app"""
        launcher, command = interpret(split('jsrun --nrs 2 --gpu_per_rs 1 ./app'))
        self.assertIn('--gpu_per_rs', launcher)
        self.assertEqual(command, ['./app'])

    def test_double_dash(self):
        """jsrun double-dash separator - interpret() strips '--' from both sides."""
        launcher, command = interpret(['jsrun', '-n', '4', '--', './app'])
        self.assertIn('jsrun', launcher)
        self.assertNotIn('./app', launcher)
        self.assertEqual(command, ['./app'])

    def test_program_with_args(self):
        """jsrun -n 2 ./app -x 1"""
        launcher, command = interpret(split('jsrun -n 2 ./app -x 1'))
        self.assertEqual(launcher, ['jsrun', '-n', '2'])
        self.assertEqual(command, ['./app', '-x', '1'])
