#!/usr/bin/expect
############################################################################
# Purpose: Test of SLURM functionality
#          Validate scontrol create, delete, and update for reservations.
#
# 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) 2009 Lawrence Livermore National Security
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Dave Bremer <dbremer@llnl.gov>
# CODE-OCEC-09-009. All rights reserved.
#
# 
# This file is part of SLURM, a resource management program.
# For details, see <https://computing.llnl.gov/linux/slurm/>.
# 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       "3.11"
set file_in       "test$test_id.input"
set exit_code     0
set res_name      ""
set user_name     ""
set def_partition ""
set def_node      ""
set ii 0

print_header $test_id


#
# Procedure to create a new reservation and validate it.
# Modifies the global var res_name in the process
#
proc create_res { res_params failure_expected } {
	#exp_internal 1
	global scontrol
	global alpha_numeric_under
	global res_name

	set ret_code 0
	set res_name ""

	#
	# Create a reservation using the list of params in res_params
	#
	set arglist [linsert $res_params 0 $scontrol create res]
	eval spawn $arglist
	expect {
		-re "Reservation created: ($alpha_numeric_under)" {
			set res_name $expect_out(1,string)
		}
		-re "Error creating the reservation: Invalid user" {
			if {!$failure_expected} {
				send_user "\nFAILURE: user not authorized "
				send_user "to create reservation\n"
			}
			set ret_code 1
			exp_continue
		}
		-re "Error" {
			if {!$failure_expected} {
				send_user "\nFAILURE: problem creating "
				send_user "reservation with args: $res_params\n"
			}
			set ret_code 1
			exp_continue
		}
		-re "error" {
			if {!$failure_expected} {
				send_user "\nFAILURE: problem creating "
				send_user "reservation with args: $res_params\n"
			}
			set ret_code 1
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			set ret_code 1
		}
		eof {
			wait
		}
	}
	if { $ret_code != 0 } {
		return $ret_code
	}

	spawn $scontrol show res $res_name
	expect {
		-re "ReservationName=($alpha_numeric_under)" {
			set tmp_res_name $expect_out(1,string)
			if {$tmp_res_name != $res_name} {
				if {!$failure_expected} {
					send_user "\nFAILURE: problem showing "
					send_user "reservation created with:  "
					send_user "$res_params\n"
				}
				set ret_code 1
				exp_continue
			}
		}
		-re "No reservations in the system" {
			if {!$failure_expected} {
				send_user "\nFAILURE: no reservations found  "
				send_user "after reservation created with:  "
				send_user "$res_params\n"
			}
			set ret_code 1
			exp_continue
		}
		-re "Reservation ($alpha_numeric_under) not found" {
			if {!$failure_expected} {
				send_user "\nFAILURE: Reservation $res_name not"
				send_user "found after reservation created "
				send_user "with:  $res_params\n"
			}
			set ret_code 1
			exp_continue
		}
	}
	#exp_internal 0

	return $ret_code
}


#
# Procedure to update a reservation
#
proc update_res { res_name res_params failure_expected } {
	global scontrol
	global alpha_numeric_under
	set ret_code 0

	#
	# Update the reservation using the list of arguments in res_params
	#
	set arglist [linsert $res_params 0 $scontrol update ReservationName=$res_name]
	eval spawn $arglist
	expect {
		-re "Reservation updated." {
			exp_continue
		}
		-re "Error creating the reservation: Invalid user" {
			if {!$failure_expected} {
				send_user "\nWARNING: user not authorized "
				send_user "to update reservation\n"
			}
			set ret_code 1
			exp_continue
		}
		-re "Error" {
			if {!$failure_expected} {
				send_user "\nFAILURE: problem updating "
				send_user "reservation $res_name with "
				send_user "$res_params\n"
			}
			set ret_code 1
			exp_continue
		}
		-re "error" {
			if {!$failure_expected} {
				send_user "\nFAILURE: problem updating "
				send_user "reservation $res_name with "
				send_user "$res_params\n"
			}
			set ret_code 1
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			set ret_code 1
		}
		eof {
			wait
		}
	}
	return $ret_code
}


#
# Procedure to delete a reservation
#
proc delete_res { res_name } {
	global scontrol
	set ret_code 0
	
	spawn $scontrol delete ReservationName=$res_name
	expect {
		-re "invalid" {
			send_user "\nFAILURE: problem deleting reservation $res_name\n"
			set ret_code 1
			exp_continue
		}
		-re "reservation is in use" {
			send_user "\nFAILURE: $res_name is in use\n"
			set ret_code 1
			exp_continue
		}
	}
	return $ret_code
}




#
# Identify usable nodes in default partition
#
spawn $sinfo -h -o %32P
expect {
	-re "($alpha_numeric_under)(\\*)" {
		set def_partition $expect_out(1,string)
		exp_continue
	}
	eof {
		wait
	}
}
if {[string compare $def_partition ""] == 0} {
	send_user "\nFAILURE: failed to find default partition\n"
	exit 1
}
spawn $sinfo -h -o "=%N=" -p $def_partition
expect {
	-re "=(.+)=" {
		set def_node $expect_out(1,string)
		exp_continue
	}
	eof {
		wait
	}
}
if {[string compare $def_node ""] == 0} {
	send_user "\nFAILURE:default partition seems to have no nodes\n"
	exit 1
}


#
# Get the user name
#
spawn $bin_id -un
expect {
	-re "($alpha_numeric_under)" {
		set user_name $expect_out(1,string)
	}
	eof {
		wait
	}
}


#
# TEST 1
# Make a list of lists with a series of parameters to test.  All the tests 
# in goodtests should pass, all those in badtests should fail.
#
set badtests "
	{}
	{Duration=5   Nodes=$def_node   User=$user_name}
	{StartTime=now   Nodes=$def_node   User=$user_name}
	{StartTime=midnight   Duration=600   User=$user_name}
	{StartTime=now   Duration=5   Nodes=ALL}
	{StartTime=now   Duration=5   NodeCnt=  Nodes=   User=$user_name}
	{StartTime=now   Duration=5   User=$user_name}
	{StartTime=blah   Duration=5   Nodes=$def_node   User=$user_name}
	{StartTime=now   Duration=foo   Nodes=$def_node   User=$user_name}
	{StartTime=now   Duration=5   Nodes=$def_node   User=$user_name  PartitionName=badpartname}
	{StartTime=now   Duration=5   Nodes=$def_node   User=$user_name  Flags=badtype}
	{StartTime=now+10minutes   EndTime=now   Nodes=$def_node   User=$user_name}
"
#	{StartTime=now   Duration=5   Nodes=$def_node   Account=badaccountname}

foreach test $badtests {
	set ret_code [create_res $test 1]
	if {$ret_code == 0} {
		send_user "\nFAILURE: Reservation $test did not fail but should have\n"
		delete_res $res_name
		exit 1
	} else {
		send_user "Expected error.  You can turn that frown upside-down.\n"
	}
}

if {[test_super_user] == 0} {
	send_user "\nWARNING: can not test more unless SlurmUser or root\n"
	exit $exit_code
}

set goodtests "
	{StartTime=now   Duration=5   Nodes=$def_node   User=$user_name}
	{StartTime=now+5minutes   EndTime=now+10minutes   Nodes=$def_node   User=$user_name}
	{StartTime=midnight   Duration=600   Nodes=$def_node   User=$user_name}
	{StartTime=now   Duration=5   Nodes=ALL   User=$user_name}
	{StartTime=now   Duration=5   NodeCnt=1   User=$user_name}
	{StartTime=now   Duration=5   Nodes=$def_node   User=$user_name  PartitionName=$def_partition}
	{StartTime=now   Duration=5   Nodes=$def_node   User=$user_name  Flags=Maint}
"
foreach test $goodtests {
	set ret_code [create_res $test 0]
	if {$ret_code != 0} {
		send_user "\nFAILURE: Unable to create a valid reservation\n"
		exit $ret_code
	}
	set ret_code [delete_res $res_name]
	if {$ret_code != 0} {
		send_user "\nFAILURE: Unable to delete a reservation\n"
		exit $ret_code
	}
}


#
# TEST 2
# Create a reservation and update it in various ways
#
set ret_code [create_res "StartTime=now+60minutes Duration=60 NodeCnt=1 User=$user_name" 0]
if {$ret_code != 0} {
	send_user "\nFAILURE: Unable to create a valid reservation\n"
	exit $ret_code
}

set goodupdates "
	{PartitionName=$def_partition}
	{PartitionName=}
	{Duration=90}
	{StartTime=now+30minutes}
	{Nodes=$def_node}
	{EndTime=now+60minutes Flags=Maint NodeCnt=1 Nodes=}
"
#	{Flags=Maint}
#	{Flags=}

if {$user_name != "root"} {
	lappend goodupdates {Users+=root} {Users-=root}
}

foreach test $goodupdates {
	set ret_code [update_res $res_name $test 0]
	if {$ret_code != 0} {
		send_user "\nFAILURE: Unable to create a valid reservation\n"
		set exit_code 1
		break
	}

}

set ret_code [delete_res $res_name]
if {$ret_code != 0} {
	send_user "\nFAILURE: Unable to delete a reservation\n"
	exit $ret_code
}


#
# TEST 3
# Make a reservation, submit a job to it, confirm that the job is paired 
# with the reservation.
#

# Make the reservation
set ret_code [create_res "StartTime=now+60minutes Duration=60 NodeCnt=1 User=$user_name" 0]
if {$ret_code != 0} {
	send_user "\nFAILURE: Unable to create a valid reservation\n"
	exit $ret_code
}

# Make the job script
exec $bin_rm -f $file_in
make_bash_script $file_in "$bin_sleep 10"

# Submit the batch job
set sbatch_pid [spawn $sbatch -N1 --reservation=$res_name $file_in]
expect {
	-re "Submitted batch job ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sbatch not responding\n"
		slow_kill $sbatch_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	send_user "\nFAILURE: batch submit failure\n"
	exit 1
}

# Show the job, make sure reservation tag is right
spawn $scontrol show job $job_id
expect {
	-re "Reservation=($alpha_numeric_under)" {
		set tmp_res_name $expect_out(1,string)
		if {$tmp_res_name != $res_name} {
			send_user "\nFAILURE: problem submitting a job to a "
			send_user "reservation.  Job $job_id is running on "
			send_user "reservation $tmp_res_name, not $res_name\n"
			set exit_code 1
			exp_continue
		}
	}
	-re "Invalid job id specified" {
		send_user "\nFAILURE: Job $job_id not found\n"
		set exit_code 1
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: scontrol not responding\n"
		set exit_code 1
	}
	eof {
		wait
	}
}

# Cancel the job
spawn $scancel -v $job_id
expect {
	-re "Invalid job_id" {
		send_user "\nFAILURE: Error cancelling the job submitted "
		send_user "to the reservation.  Job $job_id not found\n"
		set exit_code 1
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: scancel not responding\n"
		set exit_code 1
	}
	eof {
		wait
	}
}

# Delete the reservation
set ret_code [delete_res $res_name]
if {$ret_code != 0} {
	exit $ret_code
}


#
# TEST 4
# If not running as root, make a reservation restricted to root, submit a job,
# and confirm that the job is rejected.
#
if {$user_name != "root"} {
	# Make the reservation
	set ret_code [create_res "StartTime=now+60minutes Duration=60 NodeCnt=1 User=root" 0]
	if {$ret_code != 0} {
		send_user "\nFAILURE: Unable to create a reservation\n"
		set exit_code 1
	}

	# Submit the batch job
	set denied 0
	set sbatch_pid [spawn $sbatch -N1 --reservation=$res_name $file_in]
	expect {
		-re "Submitted batch job ($number)" {
			set job_id $expect_out(1,string)
			exec $scancel $job_id
			send_user "\nFAILURE: job submit should have been denied\n"
			set exit_code 1
			exp_continue
		}
		-re "Batch job submission failed: Access denied to requested reservation" {
			# Job was correctly denied
			set denied 1
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: sbatch not responding\n"
			slow_kill $sbatch_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$denied == 0} {
		send_user "\nFAILURE: Job $job_id should have been rejected "
		send_user "from reservation restricted to root.  Expected "
		send_user "rejection message not given.\n"
		set exit_code 1
	} else {
		send_user "Expected error, no worries mate.\n"
	}
	# Delete the reservation
	set ret_code [delete_res $res_name]
	if {$ret_code != 0} {
		exit $ret_code
	}
}

# Remove the temporary job script
exec $bin_rm -f $file_in

if {$exit_code == 0} {
	send_user "\nSUCCESS\n"
}
exit $exit_code
