#!/usr/bin/env expect
############################################################################
# Purpose: Test of SPANK plugin that links libslurm and queries job info
############################################################################
# Copyright (C) SchedMD LLC.
# Copyright (C) 2008-2009 Lawrence Livermore National Security.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Morris Jette <jette1@llnl.gov>
# Modified by Nathan Rini <nate@schedmd.com>
# CODE-OCEC-09-009. All rights reserved.
#
# 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 cwd             "[$bin_pwd]"
set file_in         "$test_dir/input"
set file_out        "$test_dir/output"
set file_prog       "${cwd}/$test_name.prog"
set orig_spank_conf "$test_dir/orig_conf"
set new_spank_conf  "$test_dir/new_conf"
set spank_out       "$test_dir/spank.out"
set job_id          0

if {![is_super_user]} {
	skip "This test must be run as SlurmUser"
}

set config_dir [get_conf_path]
set spank_conf_file ${config_dir}/plugstack.conf

regexp {^(\d+\.\d+)} [get_config_param SLURM_VERSION] - ctld_slurm_ver
regexp {slurm (\d+\.\d+)} [run_command_output -fail "$scontrol -V"] - loc_slurm_ver
if {$ctld_slurm_ver ne $loc_slurm_ver} {
	skip "Slurmctld ($ctld_slurm_ver) and local Slurm ($loc_slurm_ver) versions are not the same, can not continue"
}

proc cleanup {} {
	global job_id orig_spank_conf spank_conf_file bin_cp bin_rm file_prog

	cancel_job $job_id

	#
	# 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
	}
	reconfigure

	file delete ${file_prog}.so
}

#
# Build the plugin
#
log_info "slurm_dir is $slurm_dir"
if [file exists ${slurm_dir}/lib64/libslurm.so] {
	exec $bin_cc -fPIC -g -pthread -ggdb3 -Wall -shared -I${slurm_dir}/include -Wl,-rpath=${slurm_dir}/lib64 -L${slurm_dir}/lib64 -lslurm -o ${file_prog}.so ${file_prog}.c
} else {
	exec $bin_cc -fPIC -g -pthread -ggdb3 -Wall -shared -I${slurm_dir}/include -Wl,-rpath=${slurm_dir}/lib -L${slurm_dir}/lib -lslurm -o ${file_prog}.so ${file_prog}.c
}

if {[file exists $spank_conf_file]} {
	spawn $bin_cat $spank_conf_file
	expect {
		-re "test${test_id}" {
			fail "spank plugin includes vestigial test${test_id}. You probably should manually remove it from $spank_conf_file. it was probably left over from some previous test failure."
		}
		eof {
			wait
		}
	}

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

# Create the output file with lower permissions to avoid root umask issues
exec $bin_touch	"${spank_out}"
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" {
		skip "User lacks permission to update plugstack_conf file"
	}
	eof {
		wait
	}
}
reconfigure -fail

log_info "Testing sbatch.........."
make_bash_script $file_in "$srun $bin_echo IT_RAN"
set job_id [submit_job -fail "-W -N1 -t1 -o $file_out $file_in"]
wait_for_job $job_id "DONE"

# Check the output of the job
set output ""
wait_for {[regexp -all "IT_RAN" $output] == 1} {
	set output [run_command_output -fail "$bin_cat $file_out"]
}
subtest {[regexp -all "IT_RAN" $output] == 1} "Local sbatch should print IT_RAN once with srun"

# Check the output of spank plugin
set matches 0
wait_for {$matches == 4} {
	set output [run_command_output -fail "$bin_cat $spank_out"]
	set inline [regexp -inline -all "slurm_spank_task_init $re_word_str step_id=($number) job_id=($number) array_job_id=($number) array_task_id=($number)" $output]
	for {set i 0} {$i < [llength $inline]} {incr i 5} {
		set spank_step_id [lindex $inline [expr {$i + 1}]]
		set spank_job_id [lindex $inline [expr {$i + 2}]]
		set spank_array_job_id [lindex $inline [expr {$i + 3}]]
		set spank_array_task_id [lindex $inline [expr {$i + 4}]]
		if {$spank_job_id == $job_id && $spank_array_task_id == $NO_VAL && ($spank_step_id == 0 || $spank_step_id == $SLURM_BATCH_SCRIPT)} {
			incr matches
		} else {
			subfail "Values from slurm_spank_task_init are not correct: step_id: $spank_step_id job_id: $spank_job_id array_job_id: $spank_array_job_id array_task_id: $spank_array_task_id"
		}
	}
}
subtest {$matches == 4} "slurm_spank_task_init should print 4 lines with the right info"

run_command -fail -nolog "$bin_rm $spank_out"

log_info "Testing sbatch array.........."
make_bash_script $file_in "$srun $bin_echo IT_RAN"
set array_job_id [submit_job -fail "--array=1-2 -W -N1 -t1 -o $file_out $file_in"]
wait_for_job $array_job_id "DONE"

# Check the output of the job array
set output ""
wait_for {[regexp -all "IT_RAN" $output] == 2} {
	set output [run_command_output -fail "$bin_cat $file_out"]
}
subtest {[regexp -all "IT_RAN" $output] == 2} "Local sbatch array should print IT_RAN twice with srun"

# Check the output of spank plugin
set matches 0
wait_for {$matches == 10} {
	set output [run_command_output -fail "$bin_cat $spank_out"]
	set inline [regexp -inline -all "slurm_spank_task_init $re_word_str step_id=($number) job_id=($number) array_job_id=($number) array_task_id=($number)" $output]
	for {set i 0} {$i < [llength $inline]} {incr i 5} {
		set spank_step_id [lindex $inline [expr {$i + 1}]]
		set spank_job_id [lindex $inline [expr {$i + 2}]]
		set spank_array_job_id [lindex $inline [expr {$i + 3}]]
		set spank_array_task_id [lindex $inline [expr {$i + 4}]]
		if {$spank_array_job_id >= $array_job_id && ($spank_array_task_id == 1 || $spank_array_task_id == 2) && ($spank_step_id == 0 || $spank_step_id == $SLURM_BATCH_SCRIPT)} {
			incr matches
		} else {
			subfail "Values from slurm_spank_task_init are not correct: step_id: $spank_step_id job_id: $spank_job_id array_job_id: $spank_array_job_id array_task_id: $spank_array_task_id"
		}
	}
}
subtest {$matches == 10} "slurm_spank_task_init should print 10 lines with the right info" "$matches != 10"