#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm sstat functionality
#          sstat a, n, o, p, P and v options.
############################################################################
# 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 exit_code   0
set file_in1    "test$test_id.input1"
set file_in2    "test$test_id.input2"
set file_in3    "test$test_id.input3"
set test_acct   "test$test_id\_acct"
set timeout     $max_job_delay

proc cleanup {} {
	global bin_rm file_in1 file_in2 file_in3

	exec $bin_rm -f $file_in1 $file_in2 $file_in3
}

set ac		AveCPU
set ap		AvePages
set ar		AveRSS
set av		AveVMSize
set ji		JobID
set mp		MaxPages
set mpn		MaxPagesNode
set mpt		MaxPagesTask
set mr		MaxRSS
set mrn		MaxRSSNode
set mrt		MaxRSSTask
set mvs		MaxVMSize
set mvn		MaxVMSizeNode
set mvt		MaxVMSizeTask
set mc		MinCPU
set mn		MinCPUNode
set mt		MinCPUTask
set nt		NTasks

#
# Check accounting config and bail if not found.
#
if {[get_config_param "AccountingStorageType"] ne "accounting_storage/slurmdbd"} {
	skip "This test can't be run without a usable AccountStorageType"
}

if {[get_config_param "FrontendName"] ne "MISSING"} {
	skip "This test can't be run on a front-end system"
}

if {[string compare [get_admin_level] "Administrator"]} {
	skip "This test can't be run without being an Accounting administrator.\nUse: sacctmgr mod user \$USER set admin=admin"
}

#
# Identify the user and his current default account
#
set acct_name ""
set user_name [get_my_user_name]

set s_pid [spawn $sacctmgr show user $user_name]
expect {
	-re "$user_name *($re_word_str)" {
		set acct_name $expect_out(1,string)
		exp_continue
	}
	timeout {
		slow_kill $s_pid
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}

#
# Use sacctmgr to add an account
#
set aamatches 0
set sadd_pid [spawn $sacctmgr -i add account $test_acct]
expect {
	-re "Adding Account" {
		incr aamatches
		exp_continue
	}
	-re "Nothing new added" {
		log_warn "Vestigial account $test_acct found"
		incr aamatches
		exp_continue
	}
	timeout {
		slow_kill $sadd_pid
		log_error "sacctmgr add not responding"
	}
	eof {
		wait
	}
}
if {$aamatches != 1} {
	log_error "sacctmgr had a problem adding account"
}

#
# Add self to this new account
#
set sadd_pid [spawn $sacctmgr -i create user name=$user_name account=$test_acct]
expect {
	 timeout {
		slow_kill $sadd_pid
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}

log_user 0
set real_memory 0
spawn $scontrol show node
expect {
	-re "RealMemory=($number) " {
		set real_memory $expect_out(1,string)
#		exp_continue
	}
	eof {
		wait
	}
}
log_user 1
if {$real_memory < 100} {
	set job_mem_opt  "--comment=no_mem_limit"
	set step_mem_opt "-t1"
} else {
	set job_mem_opt  "--mem=100"
	set step_mem_opt "--mem=20"
}

#
# Build input script file1
#
make_bash_script $file_in1 "
  $srun $step_mem_opt --gres=craynetwork:0 --overlap sleep 60 &
  $srun $step_mem_opt --gres=craynetwork:0 --overlap sleep 60 &
  $srun $step_mem_opt --gres=craynetwork:0 --overlap sleep 60
"
#
# Spawn a job via srun using this account
# NOTE: --mem option here and in scripts above to permit running more than one
#       step per job if DefMemPerCPU or DefMemPerNode is configured.
#
set job_id3 0
spawn $sbatch -N1 -t1 -v $job_mem_opt --gres=craynetwork:0 --account=$test_acct --output=/dev/null $file_in1
expect {
	-re "Submitted batch job ($number)" {
		set job_id3 $expect_out(1,string)
		log_debug "Found JobID to be $job_id3"
		exp_continue
	}
	timeout {
		fail "sbatch not responding"
	}
	eof {
		wait
	}
}
if {$job_id3 == 0} {
	log_error "Did not get srun job_id3"
	set exit_code 1
} else {
	set matches 0
	spawn $scontrol show job $job_id3
	expect {
		 -re "Account=$test_acct" {
			incr matches
			exp_continue
		}
		timeout {
			fail "scontrol not responding"
		}
		eof {
			wait
		}
	}
	if {$matches != 1} {
		log_error "srun failed to use specified account"
		set exit_code 1
	}
}

#
# Wait for job to start running, then status it
#
if {[wait_for_job $job_id3 "RUNNING"] != 0} {
	fail "Waiting for job to start running"
}
sleep 5
################################################################
#
# Proc: sstat_job
#
# Purpose:  Pass sstat options and test
#
# Returns: Number of matches.
#
# Input: Switch options not requiring arguments
#
################################################################

proc sstat_job { soption job_id step_id } {
	global sstat number
	set debug       0
	set matches     0
	set not_support 0
	log_debug "sstat -$soption -p -j $job_id"

	if { $soption == "-allsteps" || $soption == "a" } {
		spawn $sstat -$soption -p -j $job_id
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "$job_id\\.($number)" {
				if {$debug} {log_debug "Match $expect_out(1,string)"}
				incr matches
				exp_continue
			}
			-re "$job_id\\.batch" {
				if {$debug} {log_debug "Match batch"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sstat not responding"
			}
			eof {
				wait
			}
		}
		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-noheader" || $soption == "n" } {
		spawn $sstat -$soption -p -j $job_id
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID|MaxVMSize|MaxVMSizeNode|MaxVMSizeTask|	\
				AveVMSize|MaxRSS|MaxRSSNode|MaxRSSTask|AveRSS|	\
				MaxPages|MaxPagesNode|MaxPagesTask|AvePages|	\
				MinCPU|MinCPUNode|MinCPUTask|AveCPU|NTasks" {
				if {$debug} {log_debug "Match4"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				if {$debug} {log_debug "Match5"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sstat not responding"
			}
			eof {
				wait
			}
		}
		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-parsable" || $soption == "p" } {
		spawn $sstat -$soption -p -j $job_id.$step_id
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID\\|MaxVMSize\\|MaxVMSizeNode\\|MaxVMSizeTask\\|AveVMSize\\|MaxRSS\\|" {
				if {$debug} {log_debug "Match6"}
				incr matches
				exp_continue
			}
			-re "MaxRSSNode\\|MaxRSSTask\\|AveRSS\\|MaxPages\\|MaxPagesNode\\|MaxPagesTask\\|" {
				if {$debug} {log_debug "Match7"}
				incr matches
				exp_continue
			}
			-re "AvePages\\|MinCPU\\|MinCPUNode\\|MinCPUTask\\|AveCPU\\|NTasks\\|" {
				if {$debug} {log_debug "Match8"}
				incr matches
				exp_continue
			}
			-re "$job_id\\.$step_id\\|" {
				if {$debug} {log_debug "Match9"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sstat not responding"
			}
			eof {
				wait
			}
		}
		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-parsable2" || $soption == "P" } {
		spawn $sstat -$soption -j $job_id.$step_id
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID\\|MaxVMSize\\|MaxVMSizeNode\\|MaxVMSizeTask\\|AveVMSize\\|MaxRSS\\|" {
				if {$debug} {log_debug "Match10"}
				incr matches
				exp_continue
			}
			-re "MaxRSSNode\\|MaxRSSTask\\|AveRSS\\|MaxPages\\|MaxPagesNode\\|MaxPagesTask\\|" {
				if {$debug} {log_debug "Match11"}
				incr matches
				exp_continue
			}
			-re "AvePages\\|MinCPU\\|MinCPUNode\\|MinCPUTask\\|AveCPU\\|NTasks" {
				if {$debug} {log_debug "Match12"}
				incr matches
				exp_continue
			}
			-re "$job_id\\.$step_id\\|" {
				if {$debug} {log_debug "Match13"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sstat not responding"
			}
			eof {
				wait
			}
		}
		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-verbose" || $soption == "v" } {
		spawn $sstat -$soption -p -j $job_id
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID.MaxVMSize.MaxVMSizeNode.MaxVMSizeTask.AveVMSize.MaxRSS" {
				if {$debug} {log_debug "Match14"}
				incr matches
				exp_continue
			}
			-re "MaxRSSNode.MaxRSSTask.AveRSS.MaxPages.MaxPagesNode.MaxPagesTask" {
				if {$debug} {log_debug "Match15"}
				incr matches
				exp_continue
			}
			-re "AvePages.MinCPU.MinCPUNode.MinCPUTask.AveCPU.NTasks" {
				if {$debug} {log_debug "Match16"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				if {$debug} {log_debug "Match17"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sstat not responding"
			}
			eof {
				wait
			}
		}
		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}
}

################################################################
#
# Proc: sstat_vargs
#
# Purpose:  Pass sstat options with arguments and test
#
# Returns: Number of matches.
#
# Input: Switch options with argument
#
################################################################

proc sstat_vargs { soption vargs job_id} {
	global sstat
	set debug	0
	set matches     0
	set not_support 0
	log_debug "sstat -$soption $vargs -p -j $job_id"

	if { $soption == "o" || $soption == "-format" } {
		spawn $sstat -$soption $vargs -p -j $job_id
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "AveCPU.AvePages.AveRSS.AveVMSize" {
				if {$debug} {log_debug "Match18"}
				incr matches
				exp_continue
			}
			-re "JobID.MaxPages.MaxPagesNode.MaxPagesTask" {
				if {$debug} {log_debug "Match19"}
				incr matches
				exp_continue
			}
			-re "MaxRSS.MaxRSSNode.MaxRSSTask.MaxVMSize" {
				if {$debug} {log_debug "Match20"}
				incr matches
				exp_continue
			}
			-re "MaxVMSizeNode.MaxVMSizeTask.MinCPU.MinCPUNode" {
				if {$debug} {log_debug "Match21"}
				incr matches
				exp_continue
			}
			-re "MinCPUTask.NTasks" {
				if {$debug} {log_debug "Match22"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				if {$debug} {log_debug "Match23"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sstat not responding"
			}
			eof {
				wait
			}
		}
		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}
}
################################################################
set matches [sstat_job a $job_id3 ""]
if {$matches != 4} {
	log_error "sstat -a failed ($matches != 3)"
	set exit_code 1
}

set matches [sstat_job -allsteps $job_id3 ""]
if {$matches != 4} {
	log_error "sstat --allsteps failed ($matches != 3)"
	set exit_code 1
}

set matches [sstat_job n $job_id3 ""]
if {$matches != 1} {
	log_error "sstat -n failed ($matches != 1)"
	set exit_code 1
}

set matches [sstat_job -noheader $job_id3 ""]
if {$matches != 1} {
	log_error "sstat --noheader failed ($matches != 1)"
	set exit_code 1
}

set matches [sstat_job p $job_id3 0]
if {$matches != 4} {
	log_error "sstat -p failed ($matches != 4)"
	set exit_code 1
}

set matches [sstat_job -parsable $job_id3 0]
if {$matches != 4} {
	log_error "sstat --parsable failed ($matches != 4)"
	set exit_code 1
}

set matches [sstat_job P $job_id3 0]
if {$matches != 4} {
	log_error "sstat -P failed ($matches != 4)"
	set exit_code 1
}

set matches [sstat_job -parsable2 $job_id3 0]
if {$matches != 4} {
	log_error "sstat --parsable2 failed ($matches != 4)"
	set exit_code 1
}

set matches [sstat_job v $job_id3 0]
if {$matches != 4} {
	log_error "sstat -v failed ($matches != 4)"
	set exit_code 1
}

set matches [sstat_job -verbose $job_id3 0]
if {$matches != 4} {
	log_error "sstat --verbose failed ($matches != 4)"
	set exit_code 1
}

set matches [sstat_vargs o $ac,$ap,$ar,$av,$ji,$mp,$mpn,$mpt,$mr,$mrn,$mrt,$mvs,$mvn,$mvt,$mc,$mn,$mt,$nt $job_id3]
if {$matches != 6} {
	log_error "sstat -o failed ($matches != 6)"
	set exit_code 1
}

set matches [sstat_vargs -format $ac,$ap,$ar,$av,$ji,$mp,$mpn,$mpt,$mr,$mrn,$mrt,$mvs,$mvn,$mvt,$mc,$mn,$mt,$nt $job_id3]
if {$matches != 6} {
	log_error "sstat --format failed ($matches != 6)"
	set exit_code 1
}

#
# Use sacctmgr to delete the test account (shouldn't work, because of running jobs)
#
set damatches 0
set needtodel 0
set sadel_pid [spawn $sacctmgr -i delete account $test_acct]
expect {
	-re "Deleting account" {
		incr damatches
		exp_continue
	}
	"Job(s) active, cancel job(s) before remove" {
		log_debug "This error was expected, no worries"
		set needtodel 1
	}

	timeout {
		slow_kill $sadel_pid
		fail "sacctmgr delete not responding"
	}
	eof {
		wait
	}
}

if {$needtodel == 1} {
	#
	# Use scancel to quite jobs
	#
	spawn $scancel -A $test_acct
	expect {
		eof {
			wait
		}
	}

	if {[wait_for_job $job_id3 DONE] != 0} {
		log_error "Error completing job $job_id3"
		cancel_job $job_id3
		set exit_code 1
	}

	#
	# Use sacctmgr to delete the test account
	#
	set damatches 0
	set sadel_pid [spawn $sacctmgr -i delete account $test_acct]
	expect {
		-re "Deleting account" {
			incr damatches
			exp_continue
		}

		timeout {
			slow_kill $sadel_pid
			fail "sacctmgr delete not responding"
		}
		eof {
			wait
		}
	}
}

if {$damatches != 1} {
	fail "sacctmgr had problems deleting the account"
}

if {$exit_code != 0} {
	fail "Test failed due to previous errors (\$exit_code = $exit_code)"
}
