#!/usr/bin/env expect
############################################################################
# Purpose: Test sbatch/srun/salloc path resolving
############################################################################
# Copyright (C) 2019 SchedMD LLC.
# Written by Nathan Rini <nate@schedmd.com>
#
# 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

set file_in	"bash"

# Increase timeout to account for sbatch --wait backoff
set timeout 20

proc test_sbatch { local } {
	global sbatch number file_in

	set rc -12345
	set job_id 0

	if { $local } {
		set file_in_loc "./$file_in"
		set exp_rc 0
	} else {
		set file_in_loc "$file_in"
		set exp_rc 1
	}
	spawn $sbatch -W -t1 --output=/dev/null --wrap "$file_in_loc -c /bin/false"
	expect {
		-re "Submitted batch job ($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}
		-re "error" {
			log_warn "sbatch job was not submitted"
			return 1
		}
		timeout {
			log_error "sbatch not responding"
			return 1
		}
		eof {
			lassign [wait] pid spawnid os_error_flag rc
		}
	}

	if { $job_id == 0 } {
		log_error "batch submit failure"
		return 1
	}

	if {$rc != 0 } {
		log_debug "This error is expected, not a problem"
	}

	if { $rc != $exp_rc } {
		log_error "sbatch error submitting job $job_id rc: $rc instead of $exp_rc"
		return 1
	}

	return 0
}

proc test_salloc { local } {
	global salloc number file_in

	set rc -12345
	set job_id 0

	if { $local } {
		set file_in_loc "./$file_in"
		set exp_rc 0
	} else {
		set file_in_loc "$file_in"
		set exp_rc 1
	}
	spawn $salloc -t1 $file_in_loc -c /bin/false
	expect {
		-re "salloc: Granted job allocation ($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}
		-re "error" {
			log_error "Job was not submitted"
			return 1
		}
		timeout {
			log_error "salloc not responding"
			return 1
		}
		eof {
			lassign [wait] pid spawnid os_error_flag rc
		}
	}

	if { $job_id == 0 } {
		log_error "salloc submit failure"
		return 1
	}

	if {$rc != 0 } {
		log_debug "This error is expected, not a problem"
	}

	if { $rc != $exp_rc } {
		log_error "salloc error submitting job $job_id rc: $rc instead of $exp_rc"
		return 1
	}

	return 0
}

proc test_srun { local } {
	global srun number file_in

	set rc -12345
	set job_id 0

	if { $local } {
		set file_in_loc "./$file_in"
		set exp_rc 0
	} else {
		set file_in_loc "$file_in"
		set exp_rc 1
	}
	spawn $srun -v -t1 $file_in_loc -c /bin/false
	expect {
		-re "launching ($number).($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}
		timeout {
			log_error "srun not responding"
			return 1
		}
		eof {
			lassign [wait] pid spawnid os_error_flag rc
		}
	}


	if { $job_id == 0 } {
		log_error "srun submit failure"
		return 1
	}

	if {$rc != 0 } {
		log_debug "This error is expected, not a problem"
	}

	if { $rc != $exp_rc } {
		log_error "srun error submitting job $job_id rc: $rc instead of $exp_rc"
		return 1
	}

	return 0
}

set ::env(PATH) "/bin:/usr/bin"

make_bash_script $file_in "/bin/true"

# Test tools expecting bash to be used instead of our local 'bash' script
if { [test_sbatch 0] } {
	exec $bin_rm $file_in
	fail "Failed testing sbatch"
}

if { [test_salloc 0] } {
	exec $bin_rm $file_in
	fail "Failed testing salloc"
}

if { [test_srun 0] } {
	exec $bin_rm $file_in
	fail "Failed testing srun"
}

# Test tools expecting our local 'bash' script to override the system bash
if { [test_sbatch 1] } {
	exec $bin_rm $file_in
	fail "Failed testing sbatch"
}

if { [test_salloc 1] } {
	exec $bin_rm $file_in
	fail "Failed testing salloc"
}

if { [test_srun 1] } {
	exec $bin_rm $file_in
	fail "Failed testing srun"
}

#cleanup
exec $bin_rm $file_in