#!/usr/bin/expect
############################################################################
# Purpose: Establish global state information for SLURM test suite
#
# To define site-specific state information, set the values in a file
# named 'globals.local'. Those values will override any specified here.
# for example:
#
# $ cat globals.local
# set slurm_dir "/usr/local"
# set mpicc     "/usr/local/bin/mpicc"
#
# If you want to have more than one test going at the same time for multiple
# installs you can have multiple globals.local files and set the
# SLURM_LOCAL_GLOBALS_FILE env var, and have that set to the correct
# globals.local file for your various installs.  The file can be named anything,
# not just globals.local.
#
############################################################################
# Copyright (C) 2002-2007 The Regents of the University of California.
# Copyright (C) 2008-2010 Lawrence Livermore National Security.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Morris Jette <jette1@llnl.gov>
# Additions 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://computing.llnl.gov/linux/slurm/>.
# Please also read the supplied 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.
############################################################################

global sacctmgr sacct salloc sattach sbatch sbcast scancel scontrol sinfo smap squeue sreport srun sstat strigger

# Conditional set.  Only set variable if variable does not yet exist.
proc cset {name value} {
	if {![uplevel 1 info exists $name]} {
		upvar $name tmp
		set tmp $value
	}
}

cset local_globals_file "./globals.local"

if {[info exists env(SLURM_LOCAL_GLOBALS_FILE)]} {
	set local_globals_file $env(SLURM_LOCAL_GLOBALS_FILE)
}

if [file exists $local_globals_file] {
	source $local_globals_file
}

#
# Specify the slurm install directory.
# Used to locate binaries, libraries, and header files.
#
cset slurm_dir   "/usr"
cset build_dir   "../../"
cset src_dir     "../../"
cset sacctmgr    "${slurm_dir}/bin/sacctmgr"
cset sacct       "${slurm_dir}/bin/sacct"
cset salloc      "${slurm_dir}/bin/salloc"
cset sattach     "${slurm_dir}/bin/sattach"
cset sbatch      "${slurm_dir}/bin/sbatch"
cset sbcast      "${slurm_dir}/bin/sbcast"
cset scancel     "${slurm_dir}/bin/scancel"
cset scontrol    "${slurm_dir}/bin/scontrol"
cset sinfo       "${slurm_dir}/bin/sinfo"
cset smap        "${slurm_dir}/bin/smap"
cset sprio       "${slurm_dir}/bin/sprio"
cset squeue      "${slurm_dir}/bin/squeue"
cset srun        "${slurm_dir}/bin/srun"
cset sreport     "${slurm_dir}/bin/sreport"
cset sshare      "${slurm_dir}/bin/sshare"
cset sstat       "${slurm_dir}/bin/sstat"
cset strigger    "${slurm_dir}/bin/strigger"

cset pbsnodes    "${slurm_dir}/bin/pbsnodes"
cset qdel        "${slurm_dir}/bin/qdel"
cset qstat       "${slurm_dir}/bin/qstat"
cset qsub        "${slurm_dir}/bin/qsub"

# If length of string partition is zero, use output of function
#	default_partition, otherwise use the partition explicitly
#	named in your globals.local file (or below) for poe commands
cset partition ""

# If using MPICH-2 or other version of MPI requiring pmi libary, use this
#cset mpicc	"/home/jette/mpich2-install/bin/mpicc"
#cset use_pmi	1
# OR for other versions of MPICH, use this
cset mpicc       "/usr/local/bin/mpicc"
cset use_pmi	0

# If using XCPU job launch, specify directory location as needed
cset xcpu_dir	"/mnt/xcpu"

cset poe	"/usr/bin/poe"
cset mpirun	"mpirun"
cset totalviewcli	"/usr/local/bin/totalviewcli"

# Set if using "--enable-memory-leak-debug" configuration option
cset enable_memory_leak_debug 0

# Pattern to match your shell prompt
#cset prompt {(%|#|\$|\]) *$}
cset prompt "(%|#|\\\$|]|\[^>]>) *(|\[^ ]* *)$"

#
# Specify locations of other executable files used
# Only the shell names (e.g. bin_bash) must be full pathnames
#
cset bin_awk 	"awk"
cset bin_bash   [exec which bash | tail -n 1]
cset bin_cat	"cat"
cset bin_cc	"gcc"
cset bin_chmod	"chmod"
cset bin_cmp	"cmp"
cset bin_cp	"cp"
cset bin_date	"date"
cset bin_diff	"diff"
cset bin_echo	"echo"
cset bin_env	"env"
cset bin_file	"file"
cset bin_id	"id"
cset bin_grep    "grep"

# Don't user $bin_hostname unless on a front-end system that
# doesn't fully use the slurmd, use $bin_printenv SLURMD_NODENAME
cset bin_hostname "hostname"

cset bin_kill	"kill"
cset bin_make	"make"
cset bin_od     "od"
cset bin_pkill	"pkill"
cset bin_ps	"ps"
cset bin_pwd	"pwd"
cset bin_rm	"rm"
cset bin_sed	"sed"
cset bin_sleep  "sleep"
cset bin_sort   "sort"
cset bin_touch  "touch"
cset bin_uname	"uname"
cset bin_usleep "usleep"
cset bin_wc	"wc"
cset bin_printenv "printenv"

#
# Let the commands complete without expect timing out waiting for a
# response. Single node jobs submitted to the default partition should
# be initiated within this number of seconds.
# for interactive slurm jobs: cset timeout $max_job_delay
#
cset max_job_delay 120

#
# Files must be propogated between nodes within this number of seconds.
# The delay may be due to NFS.
#
cset max_file_delay 60

#
# Desired job state must be reached within this number of seconds.
#
cset max_job_state_delay 360

#
# Specify the maximum number of tasks to use in the stress tests.
#
cset max_stress_tasks 4

#
# The error message that the "sleep" command prints when we run "sleep aaa".
#
cset sleep_error_message "(invalid time interval)|(bad character in argument)"

# Other common variables
set alpha                "\[a-zA-Z\]+"
set alpha_cap            "\[A-Z\]+"
set alpha_numeric        "\[a-zA-Z0-9\]+"
set alpha_numeric_comma  "\[a-zA-Z0-9_,\-\]+"
set alpha_numeric_under  "\[a-zA-Z0-9_\-\]+"
set alpha_under          "\[A-Z_\]+"
set end_of_line          "\[\r\n\]"
set number               "\[0-9\]+"
set float                "\[0-9\]+\\.?\[0-9\]+"
set whitespace		 "\[ \t\n\r\f\v\]+"
set alpha_numeric_nodelist "$alpha_numeric_under\\\[?\[$alpha_numeric_comma\]?\\\]?"
#
# Cache SlurmUser to check for SuperUser requests
#
cset super_user     0
cset super_user_set 0

################################################################
#
# Proc: cancel_job
#
# Purpose:  Cancel the specified job
#
# Returns: A non-zero return code indicates a failure.
#
# Input: job_id  -- The SLURM job id of a job we want to cancel.
#
################################################################

proc cancel_job { job_id } {
	global scancel bin_sleep

	if {$job_id == 0} {
		return 1
	}

	send_user "cancelling $job_id\n"
	exec $scancel -Q $job_id
	exec $bin_sleep 1
	return [wait_for_job $job_id "DONE"]
}


################################################################
#
# Proc: get_line_cnt
#
# Purpose:  Return size of the specified file
#
# Returns: Number of lines in the specified file.
#
# Input: file_name  -- Name of file to inspect.
#
################################################################
proc get_line_cnt { file_name } {
	global bin_wc number
	set lines 0
	spawn $bin_wc -l $file_name
	expect {
		-re "($number) " {
			set lines $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	return $lines
}

################################################################
#
# Proc: slow_kill
#
# Purpose:  Kill a process slowly, first trying SIGINT, pausing for
#       a second, then sending SIGKILL.
#
# Returns: A non-zero return code indicates a failure.
#
################################################################

proc slow_kill { pid } {
	global bin_kill

	catch {exec $bin_kill -INT $pid}
	catch {exec $bin_kill -INT $pid}
	sleep  1
	catch {exec $bin_kill -KILL $pid}

	return 0
}

################################################################
#
# Proc: get_my_nuid
#
# Purpose:  gets the name uid from the running user
#
# Returns: A non-zero return code indicates a failure.
#
#
################################################################

proc get_my_nuid {  } {
	global bin_id alpha alpha_numeric

	set uid -1

	log_user 0
	spawn $bin_id -nu
	expect {
		-re "($alpha_numeric|$alpha)" {
			set nuid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $nuid
}

################################################################
#
# Proc: get_my_uid
#
# Purpose:  gets the uid from the running user
#
# Returns: A non-zero return code indicates a failure.
#
#
################################################################

proc get_my_uid {  } {
	global bin_id number

	set uid -1

	log_user 0
	spawn $bin_id -u
	expect {
		-re "($number)" {
			set uid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $uid
}

################################################################
#
# Proc: get_my_gid
#
# Purpose:  gets the gid from the running user
#
# Returns: A non-zero return code indicates a failure.
#
#
################################################################

proc get_my_gid {  } {
	global bin_id number

	set gid -1

	log_user 0
	spawn $bin_id -g
	expect {
		-re "($number)" {
			set gid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $gid
}


################################################################
#
# Proc: kill_salloc
#
# Purpose:  Kill all salloc commands associated with this user.
#	Issue two SIGINT, sleep 1 and a SIGKILL
#
# Returns: A non-zero return code indicates a failure.
#
# NOTE: Use slow_kill instead of kill_salloc if you can capture
#       the process id
#
################################################################

proc kill_salloc {  } {
	global bin_id bin_pkill bin_sleep bin_usleep number

	spawn $bin_id -u
	expect {
		-re "($number)" {
			set uid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	catch {exec $bin_pkill -INT -u $uid salloc}
	catch {exec $bin_pkill -INT -u $uid salloc}
	sleep  1
	catch {exec $bin_pkill -KILL -u $uid salloc}

	return 0
}


################################################################
#
# Proc: kill_srun
#
# Purpose:  Kill all srun commands associated with this user.
#	Issue two SIGINT, sleep 1 and a SIGKILL
#
# Returns: A non-zero return code indicates a failure.
#
# NOTE: Use slow_kill instead of kill_srun if you can capture
#       the process id
#
################################################################

proc kill_srun {  } {
	global bin_id bin_pkill bin_sleep bin_usleep number

	spawn $bin_id -u
	expect {
		-re "($number)" {
			set uid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	catch {exec $bin_pkill -INT -u $uid srun}
	catch {exec $bin_pkill -INT -u $uid srun}
	sleep  1
	catch {exec $bin_pkill -KILL -u $uid srun}

	return 0
}


################################################################
#
# Proc: print_header
#
# Purpose:  Print header with test ID
#
# Input: job_id   -- The SLURM job id of a job we want to cancel.
#
################################################################

proc print_header { test_id } {

	send_user "============================================\n"
	send_user "TEST: $test_id\n"
}


################################################################
#
# Proc: wait_for_file
#
# Purpose:  Wait for the specified file to exist and have a
#           non-zero size. Note that if JobFileAppend=0 is
#           configured, a file can exist and be purged then
#           be re-created.
#
# Returns: A non-zero return code indicates a failure.
#
# Input: file_name   -- Name of the file to wait for.
#
################################################################

proc wait_for_file { file_name } {
	global bin_sleep max_file_delay

	for {set my_delay 0} {$my_delay <= $max_file_delay} {incr my_delay} {
		if {[file exists $file_name]} {
#			Add small delay for I/O buffering
			exec $bin_sleep 1
			return 0
		}
		exec $bin_sleep 1
	}
	send_user "\nFAILURE: Timeout waiting for file $file_name\n"
	return 1
}


################################################################
#
# Proc: wait_for_job
#
# Purpose:  Wait for a previously submitted SLURM job to reach
# the desired state, exponential back-off 1 to 10 seconds
#
# Returns: A non-zero return code indicates a failure.
#
# Input: job_id   -- The SLURM job id of a job we want to
#                    wait for.
#        desired_state -- The state you want the job to attain before
#                         returning.  Currently supports:
#                            DONE any terminated state
#                            RUNNING job is running
#
# NOTE: We sleep for two seconds before replying that a job is
# done to give time for I/O completion (stdout/stderr files)
#
################################################################

proc wait_for_job { job_id desired_state } {
	global scontrol max_job_state_delay

	# First verify that desired_state is supported
	switch $desired_state {
		"DONE" {}
		"RUNNING" {}
		default {
			send_user "Unsupported desired state: $desired_state\n"
			return 1
		}
	}

	set sleep_time  1
	set my_delay    0
	while 1 {
		set fd [open "|$scontrol -o show job $job_id"]
		gets $fd line
		catch {close $fd}
		if {[regexp {JobState\s*=\s*(\w+)} $line foo state] != 1} {
			set state "NOT_FOUND"
		}

		switch $state {
			"NOT_FOUND" -
			"CANCELLED" -
			"FAILED" -
			"TIMEOUT" -
			"NODE_FAIL" -
			"COMPLETED" {
				if {[string compare $desired_state "DONE"] == 0} {
					send_user "Job $job_id is DONE\n"
					sleep 2
					return 0
				}
				if {[string compare $desired_state "RUNNING"] == 0} {
					send_user "Job $job_id is $state, "
					send_user "but we wanted RUNNING\n"
				}
				return 1
			}
			"RUNNING" {
				if {[string compare $desired_state "RUNNING"] == 0} {
					send_user "Job $job_id is RUNNING\n"
					return 0
				}
				send_user "Job $job_id is in state $state, "
				send_user "desire $desired_state\n"
			}
			default {
				send_user "Job $job_id is in state $state, "
				send_user "desire $desired_state\n"
			}
		}

		if { $my_delay > $max_job_state_delay } {
			send_user "FAILURE: Timeout waiting for job state $desired_state\n"
			return 1
		}

		exec sleep $sleep_time
		set my_delay [expr $my_delay + $sleep_time]
		set sleep_time  [expr $sleep_time * 2]
		if { $sleep_time > 10 } {
			set sleep_time 10
		}
	}
}

################################################################
#
# Proc: wait_for_step
#
# Purpose: Wait for a job step to be found, exponential back-off
# 1 to 10 seconds
#
# Returns: A non-zero return code indicates a failure.
#
# Input: step_id   -- The SLURM step id of a job we want to
#                     wait for.
#
################################################################

proc wait_for_step { step_id } {
	global scontrol max_job_state_delay
	set sleep_time  1
	set my_delay    0
	while 1 {
		set fd [open "|$scontrol -o show step $step_id"]
		gets $fd line
		catch {close $fd}
		if {[regexp {Nodes=} $line foo] == 1} {
			return 0
		}
		if { $my_delay > $max_job_state_delay } {
			send_user "FAILURE: Timeout waiting for job step\n"
			return 1
		}

		send_user "Step $step_id not done yet waiting for $sleep_time seconds\n"
		exec sleep $sleep_time
		set my_delay [expr $my_delay + $sleep_time]
		set sleep_time  [expr $sleep_time * 2]
		if { $sleep_time > 10 } {
			set sleep_time 10
		}
	}
}


################################################################
#
# Proc: test_assoc_enforced
#
# Purpose: Determine if we need an association to run a job.
# This is based upon
# the value of AccountingStorageEnforce in the slurm.conf.
#
# Returns level of association enforcement, 0 if none
#
################################################################
proc test_assoc_enforced { } {
	global scontrol number

	log_user 0
	set assoc_enforced 0
	spawn $scontrol show config
	expect {
		-re "AccountingStorageEnforce *= associations" {
			set assoc_enforced 1
			exp_continue
		}
		eof {
			wait
		}
	}

	log_user 1
	return $assoc_enforced
}

################################################################
#
# Proc: test_assoc_enforced
#
# Purpose: Determine if we need an association to run a job.
# This is based upon
# the value of AccountingStorageEnforce in the slurm.conf.
#
# Returns level of association enforcement, 0 if none
#
################################################################
proc test_limits_enforced { } {
	global scontrol

	log_user 0
	set enforced 0
	spawn $scontrol show config
	expect {
		-re "AccountingStorageEnforce *= (\[a-z]+),limits" {
			set enforced 1
			exp_continue
		}
		eof {
			wait
		}
	}

	log_user 1
	return $enforced
}

################################################################
#
# Proc: test_gang
#
# Purpose: Determine if gang scheduling is configured
#
# Returns level of association enforcement, 0 if none
#
################################################################
proc test_gang { } {
	global scontrol

	log_user 0
	set gang 0
	spawn $scontrol show config
	expect {
		-re "PreemptMode *= .*GANG" {
			set gang 1
			exp_continue
		}
		eof {
			wait
		}
	}

	log_user 1
	return $gang
}

################################################################
#
# Proc: test_power_save
#
# Return 1 if power save mode is enabled, 0 otherwise
#
################################################################
proc test_power_save { } {
	global scontrol number

	log_user 0
	set suspend_time 0
	spawn $scontrol show config
	expect {
		-re "SuspendTime *= ($number)" {
			set suspend_time $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	if {$suspend_time == 0} {
		set power_save 0
	} else {
		set power_save 1
	}
	return $power_save
}

################################################################
#
# Proc: slurmd_user_root
#
# Return 1 if the SlurmdUser is root, 0 otherwise
#
################################################################
proc slurmd_user_root { } {
	global scontrol

	log_user 0
	set rc 0
	spawn $scontrol show config
	expect {
		-re "SlurmdUser *= root" {
			set rc 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $rc
}

################################################################
#
# Proc: test_topology
#
# Purpose: Determine if system is topology aware
#
# Returns level of association enforcement, 0 if none
#
################################################################
proc test_topology { } {
	global scontrol

	log_user 0
	set have_topology 1
	spawn $scontrol show config
	expect {
		-re "TopologyPlugin *= *topology/none" {
			set have_topology 0
			exp_continue
		}
		eof {
			wait
		}
	}

	log_user 1
	return $have_topology
}

################################################################
#
# Proc: test_track_wckey
#
# Purpose: Determine if we track workload characterization keys.
# This is based upon the value of TrackWCKey in the slurm.conf.
#
# Returns value of TrackWCKey
#
################################################################
proc test_track_wckey { } {
        global scontrol number

        log_user 0
        set track_wckey 0
        spawn $scontrol show config
        expect {
                -re "TrackWCKey *= ($number)" {
                        set track_wckey $expect_out(1,string)
                        exp_continue
                }
                eof {
                        wait
                }
        }

        log_user 1
        return $track_wckey
}

################################################################
# Proc: test_wiki_sched
#
# Return 1 if using sched/wiki or sched/wiki2 (Maui or Moab),
#        0 otherwise
################################################################
proc test_wiki_sched { } {
	global scontrol

	log_user 0
	set sched_wiki   0
	spawn $scontrol show config
	expect {
		-re "SchedulerType *= sched/wiki" {
			set sched_wiki 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $sched_wiki
}

################################################################
#
# Proc: test_account_storage
#
# Purpose: Determine if we are using a usable accounting storage
# package.
# This is based upon
# the value of AccountingStorageType in the slurm.conf.
#
# Returns 1 if the system is running an accounting storage type
# that is complete, 0 otherwise
#
################################################################

proc test_account_storage { } {
	global scontrol

	log_user 0
	set account_storage 0
	spawn $scontrol show config
	expect {
		-re "(accounting_storage/slurmdbd|accounting_storage/mysql|accounting_storage/pgsql)" {
			set account_storage 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $account_storage
}

################################################################
#
# Proc: test_using_slurmdbd
#
# Purpose: Since there is a lag at which the slurmdbd processes a job start
# we need to wait a bit to make sure the data has been set before proceeding.
# This is based upon
# the value of AccountingStorageType in the slurm.conf.
#
# Returns 1 if the system is running with slurmdbd, 0 otherwise
#
################################################################

proc test_using_slurmdbd { } {
	global scontrol

	log_user 0
	set account_storage 0
	spawn $scontrol show config
	expect {
		-re "(accounting_storage/slurmdbd)" {
			set account_storage 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $account_storage
}

################################################################
#
# Proc: priority_type
#
# Purpose: Use scontrol to determine the priority plugin
#
# Returns: Name of priority type
#
################################################################

proc priority_type {} {
	global scontrol

	log_user 0
	set name ""
	set fd [open "|$scontrol show config"]
	while {[gets $fd line] != -1} {
		if {[regexp {^PriorityType *= priority/(\w+)} $line frag name]
				== 1} {
			break
		}
	}
	catch {close $fd}
	log_user 1

	if {[string length $name] == 0} {
		send_user "ERROR: could not identify the Priority Type\n"
	}

	return $name
}

################################################################
#
# Proc: get_min_job_age
#
# Purpose: Use scontrol to determine the MinJobAge
#
# Returns: MinJobAge value
#
################################################################

proc get_min_job_age {} {
	global scontrol number

	set age 0
	log_user 0
	spawn $scontrol show config
	expect {
		-re "MinJobAge *= ($number)" {
			set age $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	if {$age == 0} {
		send_user "ERROR: could not identify the MinJobAge\n"
	}
	return $age
}

################################################################
#
# Proc: get_default_acct
#
# Purpose: get users default account.
#
# Returns name of default account if exists, NULL otherwise
#
################################################################

proc get_default_acct { user } {
	global sacctmgr alpha_numeric_under bin_id

	log_user 0
	set def_acct ""

	if { !$user } {
		spawn $bin_id -un
		expect {
	       	       -re "($alpha_numeric_under)" {
		       	   set user $expect_out(1,string)
		       }
		       eof {
		    	   wait
		       }
		}
	}

	spawn $sacctmgr -n list user $user format="DefaultAccount"
	expect {
		-re "($alpha_numeric_under)" {
			set def_acct $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $def_acct
}

################################################################
#
# Proc: test_front_end
#
# Purpose: Determine if the execution host is one in which the
# slurmd daemon executes on a front-end node rather than the
# compute hosts (e.g. Blue Gene systems).
#
# Returns 1 if the system uses a front-end, 0 otherwise
#
################################################################

proc test_front_end { } {
	global enable_front_end scontrol

	log_user 0
	set front_end 0
	spawn $scontrol show frontend
	expect {
		"FrontendName=" {
			set front_end 1
			exp_continue
		}
		"select/cray" {
			set front_end 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $front_end
}

################################################################
#
# Proc: test_multiple_slurmd
#
# Returns 1 if running multple slurmd per node
#
################################################################

proc test_multiple_slurmd { } {
        global scontrol

        log_user 0
        set multiple_slurmd 0
        spawn $scontrol show config
        expect {
                "MULTIPLE_SLURMD" {
                        set multiple_slurmd 1
                        exp_continue
                }
                eof {
                        wait
                }
        }
        log_user 1

        return $multiple_slurmd
}


################################################################
#
# Proc: test_xcpu
#
# Purpose: Determine if the system xcpu for task launch
#
# Returns 1 if the system is xcpu, 0 otherwise
#
################################################################

proc test_xcpu { } {
	global scontrol

	log_user 0
	set have_xcpu 0
	spawn $scontrol show config
	expect {
		"HAVE_XCPU" {
			set have_xcpu 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $have_xcpu
}

################################################################
#
# Proc: test_bluegene
#
# Purpose: Determine if the system is a bluegene system
#
# Returns 1 if the system is a bluegene, 0 otherwise
#
################################################################

proc test_bluegene { } {
	global scontrol bin_bash bin_grep

	log_user 0
	set bluegene 0
	spawn -noecho $bin_bash -c "exec $scontrol show config | $bin_grep SelectType"
	expect {
		"select/bluegene" {
			set bluegene 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $bluegene
}

################################################################
#
# Proc: test_cray
#
# Purpose: Determine if the system is a cray system
#
# Returns 1 if the system is a cray, 0 otherwise
#
################################################################

proc test_cray { } {
	global scontrol bin_bash bin_grep

	log_user 0
	set cray 0
	spawn -noecho $bin_bash -c "exec $scontrol show config | $bin_grep SelectType"
	expect {
		"select/cray" {
			set cray 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $cray
}
################################################################
#
# Proc: get_cycle_count
#
# Purpose: For tests with iteration counts (e.g. test9.1, test9.2)
#	   return the desired iteration count
#
# Returns desired iteration count
#
################################################################

proc get_cycle_count { } {
	global enable_memory_leak_debug

	if {$enable_memory_leak_debug != 0} {
		return 2
	}
	if {[test_wiki_sched] == 1} {
		return 5
	}
	return 100
}

################################################################
#
# Proc: test_select_type
#
# Purpose: Determine which select plugin is being used
#
# Returns name of select plugin
#
################################################################

proc test_select_type { } {
	global scontrol bin_bash bin_grep alpha_numeric_under

	log_user 0
	set type ""
	spawn -noecho $bin_bash -c "exec $scontrol show config | $bin_grep SelectType"
	expect {
		-re "select/($alpha_numeric_under)" {
			set type $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $type
}

################################################################
#
# Proc: test_select_type_params
#
# Purpose: Determine SelectTypeParameters being used
#
# Returns value of SelectTypeParameters
#
################################################################

proc test_select_type_params { } {
	global scontrol bin_bash bin_grep alpha_numeric_under

	log_user 0
	set params ""
	spawn -noecho $bin_bash -c "exec $scontrol show config | $bin_grep SelectTypeParameters"
	expect {
		-re "SelectTypeParameters *= *($alpha_numeric_under)" {
			set params $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1

	return $params
}

################################################################
#
# Proc: test_aix
#
# Purpose: Determine if the system is AIX
#
# Returns 1 if the system is AIX, 0 otherwise
#
################################################################

proc test_aix {} {
	set fd [open "|uname"]
	gets $fd line
	close $fd
	if {[string compare $line AIX] == 0} {
		return 1
	} else {
		return 0
	}
}

################################################################
#
# Proc: test_super_user
#
# Purpose: Determine if user is a SLURM super user (i.e. user
# root or configured SlurmUser)
#
################################################################

proc test_super_user { } {
	global alpha_numeric_under bin_id number scontrol super_user super_user_set

	if {$super_user_set != 0} {
		return $super_user
	}

#
#	Check if user root
#
	log_user 0
	spawn $bin_id -u
	set uid -1
	expect {
		-re "($number)" {
			set uid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	if {$uid == 0} {
		log_user 1
		set super_user 1
		set super_user_set 1
		return $super_user
	}

#
#	Check if SlurmUser
#
	spawn $bin_id -un
	set user ""
	expect {
		-re "($alpha_numeric_under)" {
			set user $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	spawn $scontrol show config
	set slurm_user ""
	expect {
		-re "SlurmUser *= ($alpha_numeric_under)" {
			set slurm_user $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	if {[string compare $user $slurm_user] == 0} {
		set super_user 1
	}
	set super_user_set 1
	log_user 1
	return $super_user
}

# Create a 16 bit hex number from a signed decimal number
# Replace all non-decimal characters
# Courtesy of Chris Cornish
# http://aspn.activestate.com/ASPN/Cookbook/Tcl/Recipe/415982
proc dec2hex16 {value} {
	regsub -all {[^0-x\.-]} $value {} newtemp
	set value [string trim $newtemp]
	if {$value < 32767 && $value > -32768} {
		set tempvalue [format "%#010X" [expr $value]]
		return [string range $tempvalue 6 9]
	} elseif {$value < 32768} {
		return "8000"
	} else {
		return "7FFF"
	}
}

# Create a 32 bit hex number from a signed decimal number
# Replace all non-decimal characters
# Courtesy of Chris Cornish
# http://aspn.activestate.com/ASPN/Cookbook/Tcl/Recipe/415982
proc dec2hex {value} {
	regsub -all {[^0-x\.-]} $value {} newtemp
	set value [string trim $newtemp]
	if {$value < 2147483647 && $value > -2147483648} {
		set tempvalue [format "%#010X" [expr $value]]
		return [string range $tempvalue 2 9]
	} elseif {$value < -2147483647} {
		return "80000000"
	} else {
		return "7FFFFFFF"
	}
}

################################################################
#
# Proc: available_nodes
#
# Purpose: Check to see if a given partition has a at least
#          "num_nodes" number of nodes in the alloc, idle, or comp
#          state.  This can be used to avoid launching a job that
#          will never run because nodes are in the "drained" state
#          or otherwise unavailable.
#
# Returns: Returns the number of available nodes in the partition, or
#          -1 on failure.
#
# Input: partition - name of a partition
#
################################################################

proc available_nodes { partition } {
	global sinfo

	set available -1
	send_user "$sinfo --noheader --partition $partition --state idle,alloc,comp --format %D\n"
	set fd [open "|$sinfo --noheader --partition $partition --state idle,alloc,comp --format %D"]
	gets $fd line
	catch {close $fd}
	regexp {\d+} $line available
	return $available
}


################################################################
#
# Proc: default_partition
#
# Purpose: Use scontrol to determine the name of the default partition
#
# Returns: Name of the current default partition
#
################################################################

proc default_partition {} {
	global scontrol

	set name ""
	set fd [open "|$scontrol --all --oneliner show partition"]
	while {[gets $fd line] != -1} {
		if {[regexp {^PartitionName=(\w+).*Default=YES} $line frag name]
				== 1} {
			break
		}
	}
	catch {close $fd}

	if {[string length $name] == 0} {
		send_user "ERROR: could not identify the default partition\n"
	}

	return $name
}

################################################################
#
# Proc: switch_type
#
# Purpose: Use scontrol to determine the switch type
#
# Returns: Name of SwitchType
#
################################################################

proc switch_type {} {
	global scontrol

	set name ""
	set fd [open "|$scontrol show config"]
	while {[gets $fd line] != -1} {
		if {[regexp {^SwitchType *= switch/(\w+)} $line frag name]
				== 1} {
			break
		}
	}
	catch {close $fd}

	if {[string length $name] == 0} {
		send_user "ERROR: could not identify the switch type\n"
	}

	return $name
}

################################################################
#
# Proc: make_bash_script
#
# Purpose: Create a bash script of name "script_name", and
#          make the body of the script "script_contents".
#          make_bash_script removes the file if it already exists,
#          then generates the #! line, and then dumps "script_contents"
#          to the file.  Finally, it makes certain that the script
#          is executable.
#
# Returns: Nothing.
#
# Input: script_name - file name for the bash script
#        script_contents - body of the script, not including the
#                          initial #! line.
#
################################################################

proc make_bash_script { script_name script_contents } {
	global bin_bash bin_chmod

	file delete $script_name
	set fd [open $script_name "w"]
	puts $fd "#!$bin_bash"
	puts $fd $script_contents
	close $fd
	exec $bin_chmod 700 $script_name
}

################################################################
# Given a hostname, return it's numeric suffix
################################################################
proc get_suffix { hostname } {
	set host_len [string length $hostname]
	for {set host_inx [expr $host_len-1]} {$host_inx >= 0} {incr host_inx -1} {
		set host_char [string index $hostname $host_inx]
		if {[string compare $host_char "0"] < 0} { break }
		if {[string compare $host_char "9"] > 0} { break }
	}
	incr host_inx

	if {$host_inx == $host_len} {
		send_user "\nHostname lacks a suffix:$hostname\n"
		return "-1"
	}

#	Strip off leading zeros to avoid doing octal arithmetic
	set suffix [string range $hostname $host_inx $host_len]
	set suffix_len [string length $suffix]
	for {set suffix_inx 0} {$suffix_inx < [expr $suffix_len - 1]} {incr suffix_inx} {
		set suffix_char [string index $suffix $suffix_inx]
		if {[string compare $suffix_char "0"] != 0} { break }
	}

	return [string range $suffix $suffix_inx $suffix_len]
}

################################################################
# Check if we are user root or SlurmUser
#
# Return 1 if true, 0 if false
################################################################

proc is_super_user { } {
	global alpha_numeric_under bin_id scontrol

	log_user 0
	set user_name "nobody"
	spawn $bin_id -u -n
	expect {
		-re "($alpha_numeric_under)" {
			set user_name $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	if {[string compare $user_name "root"] == 0} {
		log_user 1
		return 1
	}

	set found_user 0
	spawn $scontrol show config
	expect {
		-re "SlurmUser *= $user_name" {
			set found_user 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1
	return $found_user
}

################################################################
#
#
#
################################################################
proc check_acct_associations { } {
        global sacctmgr number alpha_numeric_under

        set rc 1
	log_user 0
	send_user "Testing Associations\n"
     	#
     	# Use sacctmgr to check associations
     	#
     	set s_pid [spawn $sacctmgr -n -p list assoc wopi wopl withd format=lft,rgt,cluster]
     	expect {
	       -re "($number)\\|($number)\\|($alpha_numeric_under)\\|" {
	       	      # Here we are checking if we have duplicates and
		      # setting up an array to check for holes later

		      set cluster $expect_out(3,string)
		      if { ![info exists c_min($cluster)] } {
			      set c_min($cluster) -1
			      set c_max($cluster) -1
		      }

	    	      set num1 $expect_out(1,string)
		      set num2 $expect_out(2,string)
		      set first [info exists found($cluster,$num1)]
		      set sec [info exists found($cluster,$num2)]
		      #send_user "$first=$num1 $sec=$num2\n";
		      if { $first } {
			     send_user "FAILURE: $cluster found lft $num1 again\n"
			     set rc 0
		      } elseif { $sec } {
			     send_user "FAILURE: $cluster found rgt $num2 again\n"
			     set rc 0
		      } else {
			     set found($cluster,$num1) 1
			     set found($cluster,$num2) 1
			     if { $c_min($cluster) == -1
				  || $c_min($cluster) > $num1 } {
			     	    set c_min($cluster) $num1
			     }
			     if { $c_max($cluster) == -1
				  || $c_max($cluster) < $num2 } {
			     	    set c_max($cluster) $num2
			     }
		      }
		      exp_continue
 	       }
	       timeout {
		      send_user "FAILURE: sacctmgr add not responding\n"
		      slow_kill $s_pid
		      set exit_code 1
	       }
	       eof {
		      wait
	       }
        }

	foreach cluster [array names c_min] {
		# Here we are checking for holes in the list from above
		for {set inx $c_min($cluster)} {$inx < $c_max($cluster)} {incr inx} {
			if { ![info exists found($cluster,$inx)] } {
				send_user "FAILURE: $cluster No index at $inx\n"
				set rc 0
			}
		}
	}
	log_user 1
	return $rc
}

################################################################
#
#
#
################################################################
proc check_accounting_admin_level { } {
	global sacctmgr alpha alpha_numeric_under bin_id

	set admin_level ""
	set user_name ""

	log_user 0

	spawn $bin_id -u -n
	expect {
		-re "($alpha_numeric_under)" {
			set user_name $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}

	if { ![string length $user_name] } {
	   	send_user "FAILURE: No name returned from id\n"
		return ""
	}

     	#
     	# Use sacctmgr to check admin_level
     	#
     	set s_pid [spawn $sacctmgr -n -P list user $user_name format=admin]
     	expect {
		-re "($alpha)" {
	    	      set admin_level $expect_out(1,string)
		      exp_continue
 	       }
	       timeout {
		      send_user "FAILURE: sacctmgr add not responding\n"
		      slow_kill $s_pid
		      set exit_code 1
	       }
	       eof {
		      wait
	       }
        }

	log_user 1
	return $admin_level
}

################################################################
#
#
#
################################################################
proc get_cluster_name { } {
	global scontrol alpha_numeric_under
	#
	# Use scontrol to find the cluster name
	#
	log_user 0
	set cluster_name ""
	set scon_pid [spawn -noecho $scontrol show config]
	expect {
		-re "ClusterName *= ($alpha_numeric_under)" {
			set cluster_name $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $scon_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}

	log_user 1
	return $cluster_name
}

################################################################
#
# Proc: get_bluegene_layout
#
# Purpose: Determine what layout mode the blugene system is running
#
# Returns name of layout mode if found, 0 otherwise
#
################################################################

proc get_bluegene_layout { } {
	global scontrol alpha_numeric_under

	log_user 0
	set layout 0
	set scon_pid [spawn -noecho $scontrol show config]
	expect {
		-re "LayoutMode *= ($alpha_numeric_under)" {
			set layout $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $scon_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	return $layout
}

################################################################
#
# Proc: get_bluegene_psets
#
# Purpose: Determine how many psets a midplane has in a bluegene system
#
# Returns num of psets, 0 if not set
#
################################################################

proc get_bluegene_psets { } {
	global scontrol number

	log_user 0
	set psets 0
	set scon_pid [spawn -noecho $scontrol show config]
	expect {
		-re "IONodesPerMP *= ($number)" {
			set psets $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $scon_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	return $psets
}

################################################################
#
# Proc: get_bluegene_type
#
# Purpose: Determine what kind of bluegene system we are running
#
# Returns 'L' for bluegene/L,
# 	  'P' for bluegene/P,
#	  'Q' for bluegene/Q,
# 	  0 if not set
#
################################################################

proc get_bluegene_type { } {
	global scontrol alpha

	log_user 0
	set type 0
	set scon_pid [spawn -noecho $scontrol show config]
	expect {
		-re "Bluegene/($alpha) configuration" {
			set type $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $scon_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	return $type
}

################################################################
#
# Proc: get_bluegene_procs_per_cnode
#
# Purpose: Determine how many cpus are on a cnode
#
# Returns count of cpus on a cnode or 0 if not set
#
################################################################

proc get_bluegene_procs_per_cnode { } {
	global scontrol number

	log_user 0
	set cpu_cnt 0
	set scon_pid [spawn -noecho $scontrol show config]
	expect {
		-re "NodeCPUCnt *= ($number)" {
			set cpu_cnt $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $scon_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	return $cpu_cnt
}

################################################################
#
# Proc: get_bluegene_cnodes_per_mp
#
# Purpose: Determine how many cnodes are in a midplane
#
# Returns count of nodes on a midplane or 0 if not set
#
################################################################

proc get_bluegene_cnodes_per_mp { } {
	global scontrol number

	log_user 0
	set node_cnt 1
	set scon_pid [spawn -noecho $scontrol show config]
	expect {
		-re "MidPlaneNodeCnt *= ($number)" {
			set node_cnt $expect_out(1,string)
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $scon_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	return $node_cnt
}

################################################################
#
# Proc: get_node_cnt
#
# Purpose: Determine how many nodes are on the system
#
# Returns count of nodes on system or 0 if unknown
#
################################################################

proc get_node_cnt { } {
	global scontrol

	log_user 0
	set node_cnt 0
	set scon_pid [spawn -noecho $scontrol show nodes]
	expect {
		-re "NodeName=" {
			incr node_cnt
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $scon_pid
			set exit_code 1
		}
		eof {
			wait
		}
	}
	log_user 1

	return $node_cnt
}

################################################################
#
# Proc: print_success
#
# Purpose:  Print success with test ID
#
# Input: test_id   -- The SLURM regression test ID.
#
################################################################

proc print_success { test_id } {

	send_user "\n"
	send_user "SUCCESS: test$test_id\n"
}
