#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Test of sinfo cpu total and allocated
#
# 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 Joseph Donaghy <donaghy1@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

set test_id     "4.12"
set file_in     "test$test_id.in"
set exit_code   0
set pnumsc      0
set pnumsi      0
set aprocsc     0
set aprocsi     0
set inode_found 0
set inode_name  ""
set smallest    1
set layout "static"
set select_type ""

proc scontrol_test { node proc_cnt } {
	global scontrol number prompt
	upvar spawn_id spawn_id

	set found 0
	set rc 0

	send "$scontrol show node $node\r"
	expect {
		-re "CPUAlloc=($number)" {
			set num_alloc $expect_out(1,string)
			set found 1
			if {$proc_cnt != $num_alloc} {
				send_user "\nFAILURE: requested $proc_cnt but got $num_alloc instead\n"
				set rc 1
			}
			exp_continue
		}
		-re $prompt {
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $mypid
			return 1
		}
		eof {
			wait
		}
	}

	if {!$found} {
		send_user "\nFAILURE: didn't get expected output from scontrol\n"
		set rc 1
	}

	return $rc
}

proc sinfo_test_1 { node proc_cnt total_procs idle_cpus } {
	global float number prompt sinfo slash
	upvar spawn_id spawn_id

	set found 0
	set rc 0

	send "$sinfo -o \"%C %A %N\" -h -n $node\r"
	# make sure we get by the sinfo command so we don't
	# think the %'s are a prompt
	expect {
		"%C %A %N" {
		}
	}
	expect {
		-re "($float)(K?)($slash)($float)(K?)($slash)($float)(K?)($slash)($float)(K?) ($number)($slash)($number) $node" {
			set found 1
			set num_alloc $expect_out(1,string)
			if {[string compare $expect_out(2,string) ""]} {
				set num_alloc [expr $num_alloc * 1024]
			}
			set num_idle $expect_out(4,string)
			if {[string compare $expect_out(5,string) ""]} {
				set num_idle [expr $num_idle * 1024]
			}
			set num_other $expect_out(7,string)
			if {[string compare $expect_out(8,string) ""]} {
				set num_other [expr $num_other * 1024]
			}
			set num_total $expect_out(10,string)
			if {[string compare $expect_out(11,string) ""]} {
				set num_total [expr $num_total * 1024]
			}

			if { $num_alloc != $proc_cnt } {
				send_user "\nFAILURE: sinfo 1 allocated cpus wrong, got $num_alloc but needed $proc_cnt\n"
				set rc 1
			} elseif { $num_idle != $idle_cpus } {
				send_user "\nFAILURE: sinfo 1 idle cpus wrong, got $num_idle but needed $idle_cpus\n"
				set rc 1
			} elseif { $num_total != $total_procs } {
				send_user "\nFAILURE: sinfo 1 total cpus wrong, got $num_total but needed $total_procs\n"
				set rc 1
			}
			exp_continue
		}
		-re $prompt {
		}
		timeout {
			send_user "\nFAILURE: sinfo not responding\n"
			slow_kill $mypid
			return 1
		}
		eof {
			wait
		}
	}

	if {!$found} {
		send_user "\nFAILURE: didn't get expected output from sinfo\n"
		set rc 1
	}

	return $rc
}

proc sinfo_test_2 { node proc_cnt total_procs } {
	global sinfo number prompt
	upvar spawn_id spawn_id

	set rc 0
	set num_alloc 0
	set num_idle 0
	set alloc_nodes 1
	set total_nodes 1

	set idle_nodes [expr $total_nodes - $alloc_nodes]

	send "$sinfo -o \"%t %D %N\" -h -n $node\r"
	# make sure we get by the sinfo command so we don't
	# think the %'s are a prompt
	expect {
		"%t %D %N" {
		}
	}
	expect {
		-re "alloc ($number)(K?) $node" {
			set num_alloc $expect_out(1,string)
			if {[string compare $expect_out(2,string) ""]} {
				set num_alloc [expr $inode_procs * 1024]
			}
			exp_continue
		}
		-re "idle ($number)(K?) $node" {
			set num_idle $expect_out(1,string)
			if {[string compare $expect_out(2,string) ""]} {
				set num_idle [expr $num_idle * 1024]
			}
			exp_continue
		}
		-re "mix ($number)(K?) $node" {
			set num_alloc $expect_out(1,string)
			if {[string compare $expect_out(2,string) ""]} {
				set num_alloc [expr $inode_procs * 1024]
			}
			exp_continue
		}
		-re $prompt {
		}
		timeout {
			send_user "\nFAILURE: sinfo not responding\n"
			slow_kill $mypid
			return 1
		}
		eof {
			wait
		}
	}

	if { $num_alloc != $alloc_nodes } {
		send_user "\nFAILURE: sinfo 2 allocated nodes wrong, got $num_alloc but needed $alloc_nodes\n"
		set rc 1
	} elseif { $num_idle != $idle_nodes } {
		send_user "\nFAILURE: sinfo 2 idle nodes wrong, got $num_idle but needed $idle_nodes\n"
		set rc 1
	}

	return $rc
}

# allocate a set of nodes (node_cnt) and the quit right after
proc allocate_and_quit { node proc_cnt total_procs } {
	global bin_bash salloc scontrol sinfo number alpha_numeric_under
	global prompt select_type procs_per_node test_id bin_rm

	set file_in "test$test_id.input"
	set job_id 0
	set num_alloc 0
	set block ""
	set rc 0
	set timeout 60
	set idle_cpus [expr $total_procs - $proc_cnt]

	make_bash_script $file_in "
	  export PS1=\"$prompt\"
	  $bin_bash --norc
	"

	set mypid [spawn $salloc -w $node -N 1 -t 5 -n $proc_cnt ./$file_in]
	expect {
		-re "Granted job allocation ($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}

		-re $prompt {
			# test for scontrol to give me the correct cpu count
			if { [scontrol_test $node $proc_cnt] } {
				send "exit\r"
				return 1
			}


			# test for sinfo to give me the correct cpu count
			if { [sinfo_test_1 $node $proc_cnt $total_procs $idle_cpus] } {
				send "exit\r"
				return 1
			}
			# test for sinfo to give me the correct node count
			if { [sinfo_test_2 $node $proc_cnt $total_procs] } {
				send "exit\r"
				return 1
			}
			send "exit\r"
			exp_continue
		}

		-re "Unable to contact" {
			send_user "\nFAILURE: slurm appears to be down\n"
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: salloc not responding\n"
			if {$job_id != 0} {
				cancel_job $job_id
			}
			slow_kill $mypid
			return 1
		}
		eof {
			wait
		}
	}
	exec $bin_rm -f $file_in

	return $rc
}

############################################################################
# test starts here
############################################################################

print_header $test_id

if {[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
}

# find the default partition
set def_part [default_partition]

# find the nodes in the default partition
set def_node_name ""
spawn $sinfo -h -o "=%N=" -p $def_part
expect {
	-re "=($alpha_numeric).($number)-($number).=" {
		set def_node_name $expect_out(1,string)
		exp_continue
	}
	eof {
		wait
	}
}
if {[string compare $def_node_name ""] == 0} {
	send_user "\nWARNING: Node name format not supported for this test\n"
	exit $exit_code
}

log_user 0
set def_hostlist ""
set part_exclusive 0
spawn $scontrol show part $def_part
expect {
	-re " Nodes=($alpha_numeric_nodelist)"  {
		set def_hostlist $expect_out(1,string)
		exp_continue
	}
	-re " BasePartitions=($alpha_numeric_nodelist)" {
		set def_hostlist $expect_out(1,string)
		exp_continue
	}
	-re " OverSubscribe=EXCLUSIVE" {
		set part_exclusive 1
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: scontrol not responding\n"
		set exit_code 1
	}
	eof {
		wait
	}
}
set host_cnt 0
spawn $scontrol show hostnames $def_hostlist
expect {
	-re "($alpha_numeric_under)"  {
		set host_name($host_cnt) $expect_out(1,string)
		incr host_cnt
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: scontrol not responding\n"
		set exit_code 1
	}
	eof {
		wait
	}
}
log_user 1
if {$host_cnt == 0} {
	send_user "\nFAILURE: could not find any nodes in default partition\n"
	exit 1
}

# find me an idle node in default partition
log_user 0
set inode_name ""
set inode_cores_per_socket 0
set inode_procs 0
set units ""
set inode_sockets 0
set inode_threads_per_core 0

set fd [open "|$scontrol --oneliner show node $def_hostlist"]
exp_internal 1
while {[gets $fd line] != -1} {
	if {[regexp {NodeName=(\w+).*CoresPerSocket=(\d+).*CPUTot=(\d+)(K?).*Sockets=(\d+).*State=IDLE ThreadsPerCore=(\d+)} $line frag inode_name inode_cores_per_socket inode_procs units inode_sockets inode_threads_per_core] == 1} {
		break
	}
	if {[regexp {NodeName=(\w+).*CoresPerSocket=(\d+).*CPUTot=(\d+)(K?).*Sockets=(\d+).*State=IDLE.CLOUD ThreadsPerCore=(\d+)} $line frag inode_name inode_cores_per_socket inode_procs units inode_sockets inode_threads_per_core] == 1} {
		break
	}
}
exp_internal 0
if {[string compare $units ""]} {
	set inode_procs [expr $inode_procs * 1024]
}
catch {close $fd}

log_user 1

if {!$inode_procs} {
	send_user "\nFAILURE: couldn't find an idle node in the default partition\n"
	exit 1
}

send_user "found idle node $inode_name with $inode_procs processors\n"

# figure out the select plugin we are using
if {[test_linear]} {
	 set smallest $inode_procs
} else {
	if {$part_exclusive == 1} {
		set smallest $inode_procs
	} elseif {[test_select_type_params "CR_CPU"]} {
		set smallest $inode_threads_per_core
	} elseif {[test_select_type_params "NONE"]} {
		set smallest $inode_threads_per_core
	} elseif {[test_select_type_params "CR_CORE"]} {
		set smallest $inode_threads_per_core
	} elseif {[test_select_type_params "CR_SOCKET"]} {
		set smallest [expr $inode_cores_per_socket *$inode_threads_per_core]
	} else {
		send_user "\nWARNING: failed to parse SelectTypeParameters '$select_params'\n"
		set smallest $inode_procs
	}
}

set exit_code [allocate_and_quit $inode_name $smallest $inode_procs]
if {!$exit_code && $smallest != $inode_procs} {
	# just to make sure we get a clean state we will sleep a bit
	sleep 1
	set exit_code [allocate_and_quit $inode_name $inode_procs $inode_procs]
}
if {$exit_code == 0} {
	send_user "\nSUCCESS\n"
} else {
	exit $exit_code
}
