#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Validate proper epilog and prolog child process management.
############################################################################
# Copyright (C) 2013 SchedMD LLC
# Written by Nathan Yee <nyee32@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 inc_slurm        "pro_epi_$test_name"
set pro_epi_prog     "$test_name\_prog"
set proc             ""
set run_host         ""
set make_file        0
set exit_code        0
set pcmd_prog        "/opt/cray/nodehealth/default/bin/pcmd"

# Use PID to avoid issues with ps aux killall
set test_pid         [pid]

if {[get_config_param "SwitchType"] eq "switch/cray"} {
	set make_file 1
	if { ![file exists $pcmd_prog] } {
		skip "The file $pcmd_prog was not found"
	}
}
if {[get_config_param "LaunchType"] ne "launch/slurm"} {
	skip "This test is only compatible with systems using launch/slurm"
}
if {![is_super_user]} {
	skip "Test can only be run as SlurmUser"
}
if {[get_config_param "FrontendName"] ne "MISSING"} {
	skip "Test can not be run with front-end configuration"
}

proc cleanup { } {
	global bin_rm test_dir config_file

	# Clean up vestigial files and restore original slurm.conf file
	exec $bin_rm -fr $test_dir
	restore_conf $config_file
	reconfigure
}

set timeout $max_job_delay

proc change_pro_epi { arg program_name conf_change task_pro } {
	global inc_slurm bin_chmod bin_rm test_dir pro_epi_prog bin_echo test_name bin_pwd re_word_str make_file bin_bash exit_code file_dir bin_ln
	global bin_sleep

	if {$conf_change == 0} {
		exec $bin_echo $arg=$test_dir/$pro_epi_prog > $test_dir/$inc_slurm
	}

	set make_ln 0
	exec $bin_rm -f $test_dir/$program_name
	spawn $bin_bash -c "$bin_ln -sv $bin_sleep $test_dir/$program_name"
	expect {
		-re "$program_name" {
			set make_ln 1
			exp_continue
		}
		timeout {
			log_error "There was a problem creating the link"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$make_ln != 1} {
		log_error "Link was not created"
		set exit_code 1
	}

	exec $bin_echo "#!$bin_bash\n" > $test_dir/$pro_epi_prog
	if {$make_file == 1} {
		exec $bin_echo "touch $file_dir/$arg.$test_name.txt" >> $test_dir/$pro_epi_prog
	}

	# task prologs are ran with fork_wait so don't do the sleep
	if {!$task_pro} {
		exec $bin_echo "$test_dir/$program_name 1000&" >> $test_dir/$pro_epi_prog
	}
	exec $bin_echo "exit 0" >> $test_dir/$pro_epi_prog

	exec $bin_chmod 700 $test_dir/$pro_epi_prog
}

proc check_file { file_name } {
	global exit_code file_dir run_host pcmd_prog

	sleep 1

	set found 0
	spawn $pcmd_prog -n $run_host "find $file_dir -name $file_name"
	expect {
		-re "$file_dir/$file_name" {
			set found 1
			exp_continue
		}
		timeout {
			fail "find is not responding"
		}
		eof {
			wait
		}
	}

	if {$found == 1} {
		log_error "The file $file_name should have been removed"
		# Root may of created this file so we can't remove it
		# exec $bin_rm -f $file_dir/$file_name
		set exit_code 1
	}
}

proc check_proc {program_name} {
	global bin_ps bin_bash bin_grep exit_code

	set proc_exist 0
	spawn $bin_bash -c "exec $bin_ps -ef | $bin_grep $program_name"
	expect {
		-re "$program_name 1000" {
			set proc_exist 1
			exp_continue
		}
		timeout {
			fail "Timed out while looking for process"
		}
		eof {
			wait
		}
	}
	if {$proc_exist == 1} {
		log_error "$program_name should have been killed"
		exec killall -9 $program_name
		set exit_code 1
	}
}

proc run_job { file_type program_name conf_change task_pro} {
	global srun run_host make_file exit_code bin_rm bin_printenv bin_sleep test_dir test_name

	# Change the included file to epilog in the slurm.conf
	change_pro_epi $file_type $program_name $conf_change $task_pro

	# Wait for NFS file propagation as needed
	$bin_sleep 10

	# Update config file
	reconfigure -fail

	set sub_job 0
	spawn $srun -t1 -w$run_host $bin_printenv SLURMD_NODENAME
	expect {
		-re "$run_host" {
			set sub_job 1
			exp_continue
		}
		timeout {
			fail "srun is not responding"
		}
		eof {
			wait
		}
	}

	if {$sub_job != 1} {
		log_error "srun did not start job"
		set exit_code 1
	}

	# wait for file
	if {$make_file == 1} {
		# Check to see that the epilog or prolog file has been removed
		check_file $file_type.$test_name.txt
	}

	if {$task_pro != 1} {
		# Check to see if the child process still exist
		check_proc $program_name
	}

	# Remove link
	exec $bin_rm -rf $test_dir/$program_name
}

#
# Copy slurm.conf file
#
set config_dir [get_conf_path]
set config_file $config_dir/slurm.conf

save_conf $config_file

set get_name 0
spawn $srun -t1 -l $bin_printenv SLURMD_NODENAME
expect {
	-re "0: *($re_word_str)" {
		set run_host $expect_out(1,string)
		set get_name 1
		exp_continue
	}
	timeout {
		log_error "Did not get run_host"
		set exit_code 1
	}
	eof {
		wait
	}
}

set file_dir "/tmp"
if {$get_name != 1} {
	fail "Did not get run_host"
}

#
# Test srun --task-epilog
#
change_pro_epi srun-task-epilog srun--task-epilog_${test_name}_${test_pid} 1 0

reconfigure -fail

set sub_job 0
set invalid_node 0
spawn $srun -t1 -w$run_host --task-epilog=$test_dir/$pro_epi_prog $bin_printenv SLURMD_NODENAME
expect {
	-re "$run_host" {
		set sub_job 1
		exp_continue
	}
	timeout {
		fail "srun is not responding"
	}
	eof {
		wait
	}
}
if {$sub_job != 1} {
	log_error "srun did not submit job"
	set exit_code 1
}

if {$make_file == 1} {
	check_file srun-task-epilog.$test_name.txt
}


# Check to see if the child process still exist
check_proc srun--task-epilog_${test_name}_${test_pid}

# Remove link
exec $bin_rm -rf $test_dir/srun--task-epilog_${test_name}_${test_pid}

#
# Task Prolog Test
#
# NOTE: This checks to see if this test is running on a Cray system. When it is
# on a Cray system it will run the user invoked srun --task-prolog and the
# slurm.conf taskprolog. Both these prologs waits for the child process to
# finish before finishing the job. In order to test the task prologs we leave
# out the child process and only check for the created file.
#

if {$make_file == 1} {

	# Test srun --task-prolog
	change_pro_epi srun-task-prolog srun--task-prolog_$test_name 1 1

	reconfigure -fail

	set sub_job 0
	set sbcast_file      "$file_dir/${test_name}_sbcast"

	spawn $salloc -t1 -w$run_host $bin_bash
	expect {
		-re "Granted job allocation ($number)" {
			send "$sbcast --force $test_dir/$pro_epi_prog $sbcast_file\n"
			send "$srun -t1 --task-prolog=$test_dir/$pro_epi_prog $bin_printenv SLURMD_NODENAME\n"
			send "exit\n"
			exp_continue
		}
		-re "$run_host" {
			set sub_job 1
			exp_continue
		}
		timeout {
			fail "salloc is not responding"
		}
		eof {
			wait
		}
	}
	if {$sub_job != 1} {
		log_error "srun did not submit job"
		set exit_code 1
	}

	check_file ${test_name}_sbcast

	check_file srun-task-prolog.$test_name.txt

	# Remove link
	exec $bin_rm -rf $test_dir/srun--task-prolog_$test_name

	# Add a line in the slurm.conf file to include the epilogs and prologs
	exec $bin_echo include $test_dir/$inc_slurm >> $config_dir/slurm.conf

	#
	# Test TaskProlog
	#
	run_job taskprolog taskprolog_${test_name}_${test_pid} 0 1

} else {
	# Add a line in the slurm.conf file to include the epilogs and prologs
	exec $bin_echo include $test_dir/$inc_slurm >> $config_dir/slurm.conf
}

###############Epilog Test###############

#
# Test Epilog param
#
run_job epilog epilog_${test_name}_${test_pid} 0 0

#
# Test TaskEpilog param
#
run_job taskepilog taskepilog_${test_name}_${test_pid} 0 0

###############Prolog Test###############

#
# Test Prolog param
#
run_job prolog prolog_${test_name}_${test_pid} 0 0

if {$exit_code} {
	fail "Test failed due to previous errors (\$exit_code = $exit_code)"
}
