#!/usr/bin/env 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://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
source ./inc3.11.1
source ./inc3.11.2
source ./inc3.11.3
source ./inc3.11.4
source ./inc3.11.5
source ./inc3.11.6
source ./inc3.11.7
source ./inc3.11.8
source ./inc3.11.9
source ./inc3.11.10
source ./inc3.11.11

set test_id		"3.11"
set file_in		"test$test_id.input"
set exit_code		0
set res_name		""
set res_name_save	""
set res_nodes		""
set res_nodecnt		0
set res_thread_cnt	0
set res_partition	""
set res_flags		""
set user_name		""
set def_oversubscribe_force 0
set def_partition	""
set def_node		""
set def_node_name	""
set def_node_inx_min	-1
set def_node_inx_max	-1
set ii			0
set log_cpu_cnt		0

print_header $test_id


proc get_res_attrs {} {

	global scontrol
	global alpha_numeric_under alpha_numeric_nodelist
	global res_name
	global res_nodes
	global res_nodecnt
	global res_corecnt
	global res_thread_cnt
	global res_partition
	global res_flags

	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
		}
		-re "Nodes=($alpha_numeric_nodelist)" {
			set res_nodes $expect_out(1,string)
			exp_continue
		}
		-re "NodeCnt=($alpha_numeric_under)" {
			set res_nodecnt $expect_out(1,string)
			exp_continue
		}
		-re "CoreCnt=($alpha_numeric_under)" {
			set res_corecnt $expect_out(1,string)
			exp_continue
		}
		-re "PartitionName=(\\S+)" {
			set res_partition $expect_out(1,string)
			exp_continue
		}
		-re "Flags=(\\S+)" {
			set res_flags $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			set ret_code 1
		}
		eof {
			wait
		}
	}
}


#
# 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 number alpha_numeric_nodelist
	global res_name
	global res_nodes
	global res_nodecnt
	global res_corecnt
	global res_thread_cnt
	global log_cpu_cnt
	global res_partition
	global res_flags

	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
		}
		-nocase "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 "\n\033\[31mFAILURE: scontrol not responding\033\[m\n"
			set ret_code 1
		}
		eof {
			wait
		}
	}
	if { $ret_code != 0 } {
		return $ret_code
	}

	get_res_attrs

	lassign [get_node_cpus $res_nodes] cpu_tot res_thread_cnt
	#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
		}
		eof {
			wait
		}
	}
	return $ret_code
}

if {[is_super_user] == 0} {
	send_user "\nWARNING: This test can't be run except as SlurmUser\n"
	exit 0
} elseif {[test_select_type_params "CR_ONE_TASK_PER_CORE"]} {
	send_user "\nWARNING: This test is incompatible SelectTypeParameters=CR_ONE_TASK_PER_CORE\n"
	exit 0
}

set have_resv 0
spawn $scontrol show res $res_name
expect {
	-re "ReservationName=($alpha_numeric_under)" {
		incr have_resv
		exp_continue
	}
	eof {
		wait
	}
}
if {$have_resv != 0} {
	send_user "*****************************************************************************************\n"
	send_user "WARNING: $have_resv reservations exist at start of test which may interfere with the test\n"
	send_user "*****************************************************************************************\n"
}

#
# Identify usable nodes in default partition
#
set def_partition [default_partition]
if {[string compare $def_partition ""] == 0} {
	send_user "\nFAILURE: failed to find default partition\n"
	exit 1
}

spawn $scontrol show partition $def_partition
expect {
	-re " OverSubscribe=FORCE" {
		set def_oversubscribe_force 1
		exp_continue
	}
	eof {
		wait
	}
}
spawn $sinfo -h -o "=%N=" -tIDLE -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
}

spawn $sinfo -h -o "=%N=" -tIDLE -p $def_partition
expect {
	-re "=($alpha_numeric).($number)-($number).=" {
		set def_node_name $expect_out(1,string)
		set def_node_inx_min $expect_out(2,string)
		set def_node_inx_max $expect_out(3,string)
		exp_continue
	}
	eof {
		wait
	}
}

# Some tests do not support names of this form "rack[0-3]_blade[0-63]"
set irregular_node_name 0
if {[string compare $def_node_name ""] == 0} {
	send_user "\nWARNING: Node name format not supported for all test\n"
	set irregular_node_name 1
}

#
# Determine system size
#
set core_spec 0
set cores_per_socket 0
set sockets_per_node 0
set heterogeneous 0
log_user 0
spawn $scontrol show node -o
expect {
	-re "CoresPerSocket=($number)" {
		if {$cores_per_socket == 0} {
			set cores_per_socket $expect_out(1,string)
		} elseif {$cores_per_socket != $expect_out(1,string)} {
			set heterogeneous 1
		}
		exp_continue
	}
	-re "Sockets=($number)" {
		if {$sockets_per_node == 0} {
			set sockets_per_node $expect_out(1,string)
		} elseif {$sockets_per_node != $expect_out(1,string)} {
			set heterogeneous 1
		}
		exp_continue
	}
	-re "CoreSpecCount=($number)" {
		if {$expect_out(1,string) > 0} {
			set core_spec 1
		}
		exp_continue
	}
	-re "CPUSpecList=($number)" {
		set core_spec 1
		exp_continue
	}
	eof {
		wait
	}
}
log_user 1
if {$heterogeneous != 0} {
	send_user "\nWARNING: Cluster is heterogeneous and some tests may fail\n"
}
set cores_per_node [expr $cores_per_socket * $sockets_per_node]
set node_count [available_nodes $def_partition ""]
set cluster_cpus [expr $cores_per_node * $node_count]

set cr_socket 0
if {[test_select_type_params "CR_SOCKET"]} {
	set cr_socket 1
}

# Retrieve a node from partition and total cores in partition
set states "idle"
set node_list [get_partition_nodes $def_partition $states]
set part_node_cnt [llength $node_list]
if {$part_node_cnt == 0} {
	send_user "\n\033\[34mWARNING: unable to get $states nodes in partition $def_partition"
	exit $exit_code
}

set part_node [lindex $node_list 0]
set part_node_cores [get_node_cores $part_node]
if {$part_node_cores == -1} {
	send_user "\n\033\[34mWARNING: unable to get cores from $part_node"
	exit $exit_code
}
set part_cores [get_part_total_cores $def_partition $states]
if {$part_cores == 0} {
	send_user "\n\033\[34mWARNING: unable to get cores from $states nodes in $def_partition"
	exit $exit_code
}

#
# Get the user name
#
set user_name [get_my_user_name]

set cons_res_actived 0
if {$def_oversubscribe_force == 0} {
	if {[test_cons_res] || [test_cons_tres]} {
		set cons_res_actived 1
	}
}
set exclusive [default_part_exclusive]
if {$cons_res_actived == 1} {
	inc3_11_1
}
inc3_11_2
inc3_11_3
inc3_11_4
inc3_11_5
inc3_11_6
if {$cons_res_actived == 1 && $part_node_cores > 1 && $core_spec == 0 && $exclusive == 0 && $cr_socket == 0} {
	inc3_11_7
	inc3_11_8
}
if {$cons_res_actived == 1 && $part_node_cores > 1 && $core_spec == 0 && $exclusive == 0 && $cr_socket == 0 && $irregular_node_name == 0 && $heterogeneous == 0} {
	inc3_11_9
}

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