#!/usr/bin/env expect
############################################################################
# Purpose: Test of SPANK plugin with pack jobs
#
# Output:  "TEST: #.#" followed by "SUCCESS" if test was successful, OR
#          "FAILURE: ..." otherwise with an explanation of the failure, OR
#          anything else indicates a failure mode that must be investigated.
############################################################################
# Copyright (C) 2017 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

set test_id	    "38.6"
set cwd		    "[$bin_pwd]"
set exit_code	    0
set file_in         "${cwd}/test${test_id}.in"
set file_out	    "${cwd}/test${test_id}.out"
set file_prog	    "${cwd}/test${test_id}.prog"
set orig_spank_conf "${cwd}/test${test_id}.orig_conf"
set new_spank_conf  "${cwd}/test${test_id}.new_conf"
set spank_out       "${cwd}/test${test_id}.spank.out"

print_header $test_id

if {[test_super_user] == 0} {
	send_user "\nWARNING: This test must be run as SlurmUser\n"
	exit 0
}
if {[test_front_end]} {
        send_user "\nWARNING: This test is incompatible with front-end systems\n"
        exit $exit_code
}
if {[test_pack_step] == 0} {
	send_user "\nWARNING: heterogeneous steps not currently supported\n"
	exit $exit_code
}

proc end_it { exit_code } {
	global test_id orig_spank_conf spank_conf_file new_spank_conf spank_out
	global file_out spank_out bin_cp bin_rm file_in file_prog
	#
	# Restore the original plugstack
	#
	if {[file exists $orig_spank_conf]} {
		exec $bin_cp -f $orig_spank_conf $spank_conf_file
	} else {
		exec $bin_rm -f $spank_conf_file
	}

       if {$exit_code == 0} {
		exec $bin_rm -f $orig_spank_conf $new_spank_conf $file_in $file_out $spank_out ${file_prog}.so
		print_success $test_id
	} else {
		send_user "\nFAILURE: This may be due to NFS sychronization problems: "
		send_user "Multiple processes on differrent nodes writing to the same file\n"
		send_user "The SlurmdLogFile on each node should include SPANK logs for each step\n"
	}
	exit $exit_code
}

set node_count [available_nodes "" ""]
if {$node_count < 2} {
        send_user "\nWARNING: Insufficient node count to run test\n"
        exit $exit_code
}

#
# Build the plugin
#
exec $bin_rm -f ${file_prog}.so
exec $bin_cc -fPIC -shared -I${slurm_dir}/include -o ${file_prog}.so ${file_prog}.c

#
# Locate slurm.conf's directory, copy the original plugstack.conf file
# and create an updated one using our new plugin
#
log_user 0
set config_dir ""
set ctld_slurm_ver ""
spawn $scontrol show config
expect {
	-re "SLURM_CONF.*= (.*)/slurm.conf.*SLURM_VERSION *= ($float)" {
		set config_dir $expect_out(1,string)
		set ctld_slurm_ver $expect_out(2,string)
		exp_continue
	}
	eof {
		wait
	}
}
log_user 1
if {[string compare $config_dir ""] == 0} {
	send_user "\nFAILURE: Could not locate slurm.conf directory\n"
	exit 1
}

log_user 0
set loc_slurm_ver ""
spawn $scontrol -V
expect {
	-re "slurm ($float)" {
		set loc_slurm_ver $expect_out(1,string)
		exp_continue
	}
	eof {
		wait
	}
}
log_user 1

if {[string compare $ctld_slurm_ver $loc_slurm_ver]} {
	send_user "\nWARNING: slurmctld ($ctld_slurm_ver) and local Slurm ($loc_slurm_ver) versions are not the same, can not continue.\n"
	exit 0
}


set spank_conf_file ${config_dir}/plugstack.conf
exec $bin_rm -f $orig_spank_conf $new_spank_conf $file_out $spank_out
if {[file exists $spank_conf_file]} {
	spawn $bin_cat $spank_conf_file
	expect {
		-re "test${test_id}" {
			send_user "\nFAILURE: spank plugin includes vestigial test${test_id}\n"
			send_user "   You probably should manually remove it from $spank_conf_file\n"
			send_user "   It was probably left over from some previous test failure.\n"
			exit 1
		}
		eof {
			wait
		}
	}

	exec $bin_cp -f $spank_conf_file $orig_spank_conf
	exec $bin_cp -f $spank_conf_file $new_spank_conf
	exec $bin_chmod 700 $spank_conf_file
	exec $bin_chmod 700 $new_spank_conf
} else {
	exec $bin_cp -f /dev/null $new_spank_conf
}

exec $bin_echo "required ${file_prog}.so ${spank_out}" >>$new_spank_conf
spawn $bin_cp -f $new_spank_conf $spank_conf_file
expect {
	-re "Permission denied" {
		send_user "\nWARNING: User lacks permission to update plugstack_conf file\n"
		exit 0
	}
	eof {
		wait
	}
}

#
# Test of locally logged messages().
# Add sleep for NFS delays in propagating $spank_conf_file
#
# NOTE: This test will fail if plugstack.conf is not in a shared location
#	The login node (where test is started) will modify the file and the
#	slurmd/slurmstepd on the compute node will not get the updated
#	configuration.
#
exec $bin_sleep 30
send_user "\n\nTest locally logged messages..........\n\n"

send_user "Checksum of local $spank_conf_file\n"
spawn $bin_sum $spank_conf_file
expect {
	-re timeout {
		send_user "\nFAILURE: sbatch not responding\n"
		end_it 1
	}
	eof {
		wait
	}
}
send_user "\n\n"

make_bash_script $file_in "
  $srun --mpi=none --test_suite_srun=5 --pack-group=0,1 $bin_sum $spank_conf_file
"

set job_id 0
set matches 0
spawn $sbatch --test_suite_sbatch=4 -t1 -o $file_out -n1 : -n1 $file_in
expect {
	-re "_test_opt_process_sbatch: opt_arg_sbatch=4" {
		incr matches
		exp_continue
	}
	-re "Submitted batch job ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	-re "slurm_spank_exit: opt_arg_sbatch=4 opt_arg_srun=0" {
		incr matches
		exp_continue
	}
	-re timeout {
		send_user "\nFAILURE: sbatch not responding\n"
		end_it 1
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	send_user "\nFAILURE: batch submit failure\n"
	end_it 1
}
if {$matches != 3} {
	send_user "\nFAILURE: spank options not processed by sbatch ($matches != 3)\n"
	end_it 1
}

#
# Wait for job to complete and check for files
#
if {[wait_for_job $job_id "DONE"] != 0} {
	send_user "\nFAILURE: waiting for job to complete\n"
	cancel_job $job_id
	end_it 1
}

# NOTE: spank logs from sbatch and srun would be intermingled here
if {[wait_for_file $file_out] == 0} {
	send_user "\n\n"
	set matches 0
	set matches_sbatch 0
	set matches_srun 0
	spawn $bin_cat $file_out
	expect {
		-re "error" {
			send_user "\nFAILURE: some error happened\n"
			end_it 1
		}
		-re "_test_opt_process_(.*?=$number)" {
			set substr $expect_out(1,string)
			if {[string equal "sbatch: opt_arg_sbatch=4" $substr]} {
				incr matches_sbatch
			} elseif {[string equal "srun: opt_arg_srun=5" $substr]} {
				incr matches_srun
			}
			exp_continue;
		}
		-re "slurm_spank_local_user_init" {
			incr matches
			exp_continue
		}
		-re "slurm_spank_exit: opt_arg_sbatch=4 opt_arg_srun=5" {
			incr matches
			exp_continue
		}
		eof {
			wait
		}
	}
	set fail_cnt 0
	if {$matches_sbatch != 2} {
		send_user "\nFAILURE: local (sbatch) sbatch spank plugin failure ($matches_sbatch != 2)\n"
		set fail_cnt 1
	}
	if {$matches_srun != 2} {
		send_user "\nFAILURE: local (srun) srun spank plugin failure ($matches_srun != 2)\n"
		set fail_cnt 1
	}
	if {$matches != 2} {
		send_user "\nFAILURE: local (srun) spank plugin failure ($matches != 2)\n"
		set fail_cnt 1
	}
	if {$fail_cnt != 0} {
		end_it 1
	}
	send_user "\nlocal (srun) spank plugin success\n"
} else {
	end_it 1
}

#
# Test for slurmd output in file
#
if {[wait_for_file $spank_out] == 0} {
	send_user "\n\n"
	set matches        0
	set matches_sbatch 0
	set matches_srun   0
	set matches_spank_exit 0
	set matches_spank_init 0
	spawn $bin_sort $spank_out
	expect {
		-re "slurm_spank_(\\S+): opt_arg_sbatch=($number) opt_arg_srun=($number)" {
			set spank_type $expect_out(1,string)
			set sbatch_arg $expect_out(2,string)
			set srun_arg   $expect_out(3,string)

			if {($spank_type eq "exit") && ($sbatch_arg eq "4")} {
#				Skip (possible) external job containers
				incr matches_spank_exit
			} elseif {$spank_type eq "task_init"} {
				incr matches_spank_init
			}

			if {($sbatch_arg eq "4") && ($srun_arg eq "0")} {
				incr matches_sbatch
			} elseif {($sbatch_arg eq "4") && ($srun_arg eq "5")} {
				incr matches_srun
			}
			exp_continue
		}
		-re "spank_get_item: argv" {
			incr matches
			exp_continue
		}
		-re "spank_get_item: my_uid=" {
			incr matches
			exp_continue
		}
		eof {
			wait
		}
	}

	set fail_cnt 0
	if {$matches_spank_exit != 3} {
		send_user "\nFAILURE: remote (slurmd - spank_exit) spank plugin failure ($matches_spank_exit != 3)\n"
		set fail_cnt 1
	}
	if {$matches_spank_init != 3} {
		send_user "\nFAILURE: remote (slurmd - spank_init) spank plugin failure ($matches_spank_init != 3)\n"
		set fail_cnt 1
	}
	if {$matches_sbatch != 2} {
		send_user "\nFAILURE: remote (slurmd) sbatch spank plugin failure ($matches_sbatch != 2)\n"
		set fail_cnt 1
	}
	if {$matches_srun != 4} {
		send_user "\nFAILURE: remote (slurmd) srun spank plugin failure ($matches_srun != 4)\n"
		set fail_cnt 1
	}
	if {$matches != 8} {
		send_user "\nFAILURE: remote (slurmd) spank plugin failure ($matches != 8)\n"
		set fail_cnt 1
	}
	if {$fail_cnt != 0} {
		send_user "\nWARNING: Check for matching checksums on the plugstack.conf file.\n"
		send_user "         Different checksums could indicate file system delays\n"
		end_it 1
	}
	send_user "\nRemote (slurmd) spank plugin success\n"
} else {
	end_it 1
}

end_it $exit_code
