#!/usr/bin/expect
############################################################################
# Purpose: Test of SLURM functionality
#          Test that the shared option in partitions is being enforced.
#
# 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) 2014 SchedMD LLC
# Written by Nathan Yee <nyee32@schedmd.com>
#
# This file is part of SLURM, a resource management program.
# For details, see <http://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       17.36
set file_in       "test$test_id\_sc"
set test_part_1   "test$test_id\_part_1"
set test_part_2   "test$test_id\_part_2"
set test_part_3   "test$test_id\_part_3"
set num_jobs      0
set cr_core       0
set cr_cpu        0
set cpu_cnt       0
set socket_cnt    0
set thread_cnt    0
set shared_j_cnt  2
set job_id        0
set node_name     ""
set exit_code     0

print_header $test_id

if {[test_select_type] != "cons_res"} {
	send_user "\nWARNING: test is only compatible with a config of SelectType=select/cons_res\n"
	exit $exit_code
}

proc cr_core_cpu { node } {

	global cr_cpu cr_core core_cnt socket_cnt scontrol number exit_code
	global cpu_cnt thread_cnt

	# Determine if CR_CPU or CR_CORE is used
	log_user 0
	spawn $scontrol show config
	expect {
		-re "SelectTypeParameters *= .*CR_CPU" {
			set cr_cpu 1
			exp_continue
		}
		-re "SelectTypeParameters *= .*CR_CORE" {
			set cr_core 1
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol is not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	spawn $scontrol show node $node
	expect {
		-re "CoresPerSocket=($number)" {
			set core_cnt $expect_out(1,string)
			exp_continue
		}
		-re "CPUTot=($number)" {
			set cpu_cnt $expect_out(1,string)
			exp_continue
		}
		-re "Sockets=($number)" {
			set socket_cnt $expect_out(1,string)
			exp_continue
		}
		-re "ThreadsPerCore=($number)" {
			set thread_cnt $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol is not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	if { $cr_cpu == 1 } {
		return [expr $cpu_cnt / $thread_cnt]
	} elseif { $cr_core == 1 } {
		return [expr $core_cnt * $socket_cnt]
	} else {
		send_user "\nFAILURE: an error occurred when getting the cpu or"
		send_user " core count\n"
		set exit_code 1
	}

}

proc create_part {part shared node} {

	global scontrol exit_code

	spawn $scontrol create partitionname=$part shared=$shared nodes=$node
	expect {
		timeout {
			send_user "\nFAILURE: scontrol is not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

}

proc delete_part {part} {
	global scontrol exit_code

	spawn $scontrol delete partitionname=$part
	expect {
		timeout {
			send_user "\nFAILURE: scontrol is not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}
}

proc sub_job { job shared part } {
	global sbatch file_in number job_id exit_code

	set job_id 0
	if { $shared } {
		spawn $sbatch -a$job -t1 -p$part --share -o/dev/null $file_in
		expect {
			-re "Submitted batch job ($number)" {
				set job_id $expect_out(1,string)
				exp_continue
			}
			timeout {
				send_user "\nFAILURE: sbatch is not responding\n"
				set exit_code 1
			}
			eof {
				wait
			}
		}
	} else {
		spawn $sbatch -a$job -t1 -p$part -o/dev/null $file_in
		expect {
			-re "Submitted batch job ($number)" {
				set job_id $expect_out(1,string)
				exp_continue
			}
			timeout {
				send_user "\nFAILURE: sbatch is not responding\n"
				set exit_code 1
			}
			eof {
				wait
			}
		}
	}
	if {$job_id == 0} {
		send_user "\nFAILURE: sbatch did not submit job\n"
		set exit_code 1
	}

}

proc check_job { exp_num_jobs } {

	global squeue job_id num_jobs file_in exit_code

	# Wait a bit for the job to start
	wait_for_job ${job_id}_0 RUNNING

	set job_cnt 0
	# If gang scheduling is configured, some jobs will be suspended
	spawn $squeue --job=$job_id --state=RUNNING,SUSPENDED -Oname --noheader
	expect {
		-re "$file_in" {
			incr job_cnt 1
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: squeue is not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if { $job_cnt != $exp_num_jobs } {
		send_user "\nFAILURE: the number of possible jobs that could "
		send_user "run were not reached ($job_cnt != $exp_num_jobs).\n"
		send_user "This could be due to memory or other limits.\n"
		set exit_code 1
	}
}

proc check_part { } {
	global scontrol test_part_1 test_part_2 test_part_3 exit_code

	set found     0
	set part_name ""
	log_user 0
	spawn $scontrol -a show part
	expect {
		-re "PartitionName=($test_part_1|$test_part_2|$test_part_3)" {
			set found 1
			set part_name $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	if {$found == 1} {
		send_user "\WARNING: partition $part_name already exists\n"
		delete_part $test_part_1
		delete_part $test_part_2
		delete_part $test_part_3
	}
}

##################################Test Begins##################################
# Check that partition names have not been used
check_part

# Remove any vestigial paritions or scripts
exec $bin_rm -f $file_in

make_bash_script $file_in "sleep 10"

spawn $sbatch -N1 -t1 -o/dev/null --exclusive $file_in
expect {
	-re "Submitted batch job ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sbatch is not reponding\n"
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	send_user "\nFAILURE: sbatch did not submit job\n"
	exit $exit_code
}

# Wait a bit for job to start
wait_for_job $job_id RUNNING

# Identify node to use for testing
set got_node 0
spawn $scontrol show job $job_id
expect {
	-re "NodeList=($alpha_numeric_under)" {
		set node_name $expect_out(1,string)
		set got_node 1
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: scontrol is not responding\n"
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$got_node != 1} {
	send_user "\nFAILURE: no node was set for job\n"
	exit 1
}

cancel_job $job_id

####################################
#
# Test partition with shared=NO
#
####################################
send_user "\n\nTest partition with shared=NO\n"

# Determine the number of cores or CPUs
set num_jobs [cr_core_cpu $node_name]

# Create first test partition
create_part $test_part_1 NO $node_name

# Submit a job with shared
sub_job "0-$num_jobs" 1 $test_part_1

# Check that the correct number of jobs are running
check_job $num_jobs
cancel_job $job_id

# Submit a job without shared
sub_job "0-$num_jobs" 0 $test_part_1
check_job $num_jobs
cancel_job $job_id
delete_part $test_part_1

####################################
#
# Test partition with shared=YES:2
#
####################################
send_user "\n\nTest partition with shared=YES:2\n"

# Make a new partition with shared=yes:2
create_part $test_part_2 "YES:$shared_j_cnt" $node_name

# Submit a job with shared (expect 2 jobs per core/CPU)
set new_job_limit [expr $num_jobs *2]
sub_job "0-$new_job_limit" 1 $test_part_2
check_job $new_job_limit
cancel_job $job_id

# Submit a job without shared (expect 1 job per core/CPU)
sub_job "0-$num_jobs" 0 $test_part_2
check_job $num_jobs
cancel_job $job_id
delete_part $test_part_2

########################################
#
# Test partition with shared=EXCLUSIVE
#
########################################
send_user "\n\nTest partition with shared=EXCLUSIVE\n"

# Make a new partition with shared=EXCLUSIVE
create_part $test_part_3 "EXCLUSIVE" $node_name

# Submit a job with shared (expected 1 job per node)
sub_job "0-$num_jobs" 1 $test_part_3
check_job 1
cancel_job $job_id

# Submit a job with shared (expected 1 job per node)
sub_job "0-$num_jobs" 0 $test_part_3
check_job 1
cancel_job $job_id
delete_part $test_part_3

if {$exit_code == 0} {
	exec $bin_rm -f $file_in
	send_user "\nSUCCESS\n"
} else {
	send_user "\nFAILURE\n"
}
exit $exit_code
