#!/usr/bin/env expect
############################################################################
# Purpose: Validate sbatch hetjob options.
#
# Reqs:    1. Using slurmdbd accounting storage type and is up
#          2. controllers are up and running.
############################################################################
# Copyright (C) SchedMD LLC.
#
# This file is part of Slurm, a resource management program.
# For details, see <https://slurm.schedmd.com/>.
# Please also read the included file: DISCLAIMER.
#
# Slurm is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option)
# any later version.
#
# Slurm is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along
# with Slurm; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA.
############################################################################

source ./globals
source ./globals_het_jobs

set file_in      "$test_dir/input"
set job_id       0

if {[get_config_param "SchedulerType"] ne "sched/backfill"} {
	skip "This test requires SchedulerType = sched/backfill"
}

proc sbatch {} {
	global number sbatch file_in

	set matches 0
	set job_id 0

	spawn $sbatch -t1 -o /dev/null $file_in
	expect {
		-re "Batch job submission failed" {
			skip "Unable to execute test due to system configuration"
		}
		-re "Submitted batch job ($number)" {
			incr matches
			set job_id $expect_out(1,string)
		}
		timeout {
			fail "sbatch not responding"
		}
		eof {
			wait
		}
	}
	if {$matches != 1} {
		fail "Batch submit failure"
	}

	return $job_id
}

proc cancel_job_id { } {
	global scancel job_id

	if {$job_id != 0} {
		exec $scancel $job_id
	}
}

proc cleanup { } {
	global bin_rm test_name

	cancel_job_id
	exec $bin_rm -f $test_name*.out
}

# Start test

make_bash_script $file_in "#SBATCH --cpus-per-task=4 --mem-per-cpu=10 --ntasks=1
#SBATCH hetjob
#SBATCH --cpus-per-task=2 --mem-per-cpu=2  --ntasks=1 -t1
#SBATCH hetjob
#SBATCH --cpus-per-task=1 --mem-per-cpu=6  --ntasks=1 -t1

$bin_sleep 300"

log_info "\n################################################################\n"
log_info "Submit hetjob and verify output from scontrol show job"
log_info "\n################################################################\n"

set job_id [sbatch]

set matches 0
set id_set  0
set id_regex "\[0-9,-\]+"

# scontrol will show all components of the hetjob, but we just look at the
# first set (hetjob leader).
spawn $scontrol show job $job_id
expect {
	-re "JobId=$job_id\\s+HetJobId=$job_id" {
		incr matches
		exp_continue
	}
	-re "HetJobOffset=0" {
		incr matches
		exp_continue
	}
	-re "HetJobIdSet=($id_regex)" {
		incr matches
		if {$id_set == 0} {
			set id_set $expect_out(1,string)
		}
		exp_continue
	}
	-re "CPUs/Task=1" {
		incr matches
		exp_continue
	}
	-re "MinMemoryCPU=6M" {
		incr matches
		exp_continue
	}
	timeout {
		fail "scontrol not responding"
	}
	eof {
		wait
	}
}
if {$matches != 7 } {
	fail "Problem with scontrol ($matches != 7)"
}

set js [parse_id_set $id_set $job_id]

set j2 [lindex $js 0]
set j3 [lindex $js 1]

set matches 0

spawn $scontrol show job $j2
expect {
	-re "JobId=$j2\\s+HetJobId=$job_id" {
		incr matches
		exp_continue
	}
	-re "HetJobOffset=1" {
		incr matches
		exp_continue
	}
	-re "HetJobIdSet=$id_regex" {
		incr matches
		exp_continue
	}
	-re "CPUs/Task=2" {
		incr matches
		exp_continue
	}
	-re "MinMemoryCPU=2M" {
		incr matches
		exp_continue
	}
	timeout {
		fail "scontrol not responding"
	}
	eof {
		wait
	}
}
if {$matches != 5} {
	fail "Problem with scontrol ($matches != 5)"
}

set matches 0

spawn $scontrol show job $j3
expect {
	-re "JobId=$j3\\s+HetJobId=$job_id" {
		incr matches
		exp_continue
	}
	-re "HetJobOffset=2" {
		incr matches
		exp_continue
	}
	-re "HetJobIdSet=$id_regex" {
		incr matches
		exp_continue
	}
	-re "CPUs/Task=1" {
		incr matches
		exp_continue
	}
	-re "MinMemoryCPU=6M" {
		incr matches
		exp_continue
	}
	timeout {
		fail "scontrol not responding"
	}
	eof {
		wait
	}
}
if {$matches != 5} {
	fail "Problem with scontrol ($matches != 5)"
}