#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Test of sinfo cpu total and allocated
############################################################################
# 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 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 ""
set job_id      0

proc cleanup {} {
	global job_id

	cancel_job $job_id
}

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

	set found 0

	send "$scontrol show node $node\r"
	expect {
		-re "CPUAlloc=($number)" {
			set num_alloc $expect_out(1,string)
			set found 1
			subtest {$proc_cnt == $num_alloc} "Verify CPUAlloc" "Requested $proc_cnt but got $num_alloc instead"
			exp_continue
		}
		-re $test_prompt {
		}
		timeout {
			fail "scontrol not responding"
		}
		eof {
			wait
		}
	}

	subtest {$found} "Verify output from scontrol"
}

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

	set found 0

	send "$sinfo -o \"%C %A %N\" -h -n $node\r"
	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 {$expect_out(2,string) ne ""} {
				set num_alloc [expr $num_alloc * 1024]
			}
			set num_idle $expect_out(4,string)
			if {$expect_out(5,string) ne ""} {
				set num_idle [expr $num_idle * 1024]
			}
			set num_other $expect_out(7,string)
			if {$expect_out(8,string) ne ""} {
				set num_other [expr $num_other * 1024]
			}
			set num_total $expect_out(10,string)
			if {$expect_out(11,string) ne ""} {
				set num_total [expr $num_total * 1024]
			}

			subtest {$num_alloc == $proc_cnt} "Verify sinfo allocated cpus" "sinfo 1 allocated cpus wrong, got $num_alloc but needed $proc_cnt"
			subtest {$num_idle == $idle_cpus} "Verify sinfo idle cpus" "sinfo 1 idle cpus wrong, got $num_idle but needed $idle_cpus"
			subtest {$num_total == $total_procs} "Verify sinfo total cpus" "sinfo 1 total cpus wrong, got $num_total but needed $total_procs"

			exp_continue
		}
		-re $test_prompt {
		}
		timeout {
			fail "sinfo not responding"
		}
		eof {
			wait
		}
	}

	subtest {$found} "Verify output from sinfo"
}

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

	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"
	expect {
		-re "alloc ($number)(K?) $node" {
			set num_alloc $expect_out(1,string)
			if {$expect_out(2,string) ne ""} {
				set num_alloc [expr $inode_procs * 1024]
			}
			exp_continue
		}
		-re "idle ($number)(K?) $node" {
			set num_idle $expect_out(1,string)
			if {$expect_out(2,string) ne ""} {
				set num_idle [expr $num_idle * 1024]
			}
			exp_continue
		}
		-re "mix ($number)(K?) $node" {
			set num_alloc $expect_out(1,string)
			if {$expect_out(2,string) ne ""} {
				set num_alloc [expr $inode_procs * 1024]
			}
			exp_continue
		}
		-re $test_prompt {
		}
		timeout {
			fail "sinfo not responding"
		}
		eof {
			wait
		}
	}

	subtest {$num_alloc == $alloc_nodes} "Verify sinfo allocated nodes" "sinfo 2 allocated nodes wrong, got $num_alloc but needed $alloc_nodes"
	subtest {$num_idle == $idle_nodes} "Verify sinfo idle nodes" "sinfo 2 idle nodes wrong, got $num_idle but needed $idle_nodes"
}

# 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 re_word_str
	global test_prompt select_type procs_per_node
	global job_id

	set num_alloc 0
	set block ""
	set timeout 60
	set idle_cpus [expr $total_procs - $proc_cnt]

	spawn $salloc -w $node -N 1 -t 5 -n $proc_cnt $bin_bash
	expect {
		-re "Granted job allocation ($number)" {
			set job_id $expect_out(1,string)
			reset_bash_prompt
			exp_continue
		}

		-re $test_prompt {
			# test for scontrol to give me the correct cpu count
			scontrol_test $node $proc_cnt

			# test for sinfo to give me the correct cpu count
			sinfo_test_1 $node $proc_cnt $total_procs $idle_cpus

			# test for sinfo to give me the correct node count
			sinfo_test_2 $node $proc_cnt $total_procs

			send "exit\r"
			exp_continue
		}

		-re "Unable to contact" {
			fail "Slurm appears to be down"
		}
		timeout {
			fail "salloc not responding"
		}
		eof {
			wait
		}
	}
}

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

set select_type_parameters [get_config_param "SelectTypeParameters"]
if {[param_contains $select_type_parameters "CR_ONE_TASK_PER_CORE"]} {
	skip "This test is incompatible SelectTypeParameters=CR_ONE_TASK_PER_CORE"
}

# Check for OverSubscribe
set part_exclusive 0
if {[get_partition_param [default_partition] "OverSubscribe"] == "EXCLUSIVE"} {
	set part_exclusive 1
}

# Get an idle node and their main params
set nodes [get_nodes_by_state]
if {[llength $nodes] == 0} {
	skip "Couldn't find an idle node in the default partition"
}

set inode_name [lindex $nodes 0]
set inode_cores_per_socket [get_node_param $inode_name "CoresPerSocket"]
set inode_procs [get_node_param $inode_name "CPUTot"]
set inode_threads_per_core [get_node_param $inode_name "ThreadsPerCore"]

log_debug "Found idle node $inode_name with $inode_procs processors"

# figure out the select plugin we are using
if {[check_config_select "linear"]} {
	 set smallest $inode_procs
} else {
	if {$part_exclusive == 1} {
		set smallest $inode_procs
	} elseif [param_contains $select_type_parameters "CR_CPU*"] {
		set smallest $inode_threads_per_core
	} elseif [param_contains $select_type_parameters "NONE"] {
		set smallest $inode_threads_per_core
	} elseif [param_contains $select_type_parameters "CR_CORE*"] {
		set smallest $inode_threads_per_core
	} elseif [param_contains $select_type_parameters "CR_SOCKET*"] {
		set smallest [expr $inode_cores_per_socket *$inode_threads_per_core]
	} else {
		log_warn "Failed to parse SelectTypeParameters '$select_params'"
		set smallest $inode_procs
	}
}

allocate_and_quit $inode_name $smallest $inode_procs
if {$smallest != $inode_procs} {
	# just to make sure we get a clean state we will sleep a bit
	sleep 1
	allocate_and_quit $inode_name $inode_procs $inode_procs
}