#!/usr/bin/expect
############################################################################
# Purpose: Test of SLURM sstat functionality
#          sstat a, n, o, p, P and v options.
#
# 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://computing.llnl.gov/linux/slurm/>.
# 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     "23.3"
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_acct"
set timeout 60
print_header $test_id

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        
set sc		SystemCPU     
set tc		TotalCPU
#
# Check accounting config and bail if not found.
#
if { [test_account_storage] == 0 } {
	send_user "\nWARNING: This test can't be run without a usable AccountStorageType\n"
	exit 0
}

if { [string compare [check_accounting_admin_level] "Administrator"] } {
	send_user "\nWARNING: This test can't be run without being an Accounting administrator.\nUse sacctmgr mod user \$USER_NAME set admin=admin.\n"
	exit 0
}

#
# Identify the user and his current default account
#
set acct_name ""
set user_name ""
set user_gid ""
spawn $bin_id -u -n
expect {
	 -re "($alpha_numeric_under)" {
		set user_name $expect_out(1,string)
		exp_continue
	}
	eof {
		wait
	}
}

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

set s_pid [spawn $sacctmgr show user $user_name]
expect {
	-re "$user_name *($alpha_numeric_under)" {
		set acct_name $expect_out(1,string)
		exp_continue
	}
	timeout {
		send_user "FAILURE: sacctmgr add not responding\n"
		slow_kill $s_pid
		exit 1
	}
	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" {
		send_user "\nWARNING: vestigial account $test_acct found\n"
		incr aamatches
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sacctmgr add not responding\n"
		slow_kill $sadd_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$aamatches != 1} {
	send_user "\nFAILURE:  sacctmgr had a problem adding account.\n"
	exit 1
}

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

#
# Build input script file1
#
make_bash_script $file_in1 "
$srun $file_in2

"

#
# Build input script file2
#
make_bash_script $file_in2 "
$srun $file_in3

"

#
# Build input script file3
#
make_bash_script $file_in3 "
$srun sleep 10

"

#
# Spawn a job via srun using this account
#
set job_id3 0
spawn $sbatch -N1 -v --account=$test_acct $file_in1
expect {
	-re "Submitted batch job ($number)" {
		set job_id3 $expect_out(1,string)
		send_user "\nFOUND JobID to be $job_id3\n"
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sbatch not responding\n"
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$job_id3 == 0} {
	send_user "\nFAILURE: did not get srun job_id3\n"
	set exit_code 1
} else {
	set matches 0
	spawn $scontrol show job $job_id3
	expect {
		 -re "Account=$test_acct" {
			incr matches
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}
	if {$matches != 1} {
		send_user "\nFAILURE: srun failed to use specified account\n"
		set exit_code 1
	}
}

#
# Wait for job to start running, then signal it
#
if {[wait_for_job $job_id3 "RUNNING"] != 0} {
	send_user "\nFAILURE: waiting for job to start running\n"
	exit 1
}
################################################################
#
# 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} {
	global sstat
	set debug       0
	set exit_code   0
	set matches     0
	set not_support 0
	send_user "sstat -$soption -p -j $job_id\n"

	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.0" {
			if {$debug} {send_user "\nmatch1\n"}
			incr matches
			exp_continue
		}
		-re "$job_id.1" {
			if {$debug} {send_user "\nmatch2\n"}
			incr matches
			exp_continue
		}
		-re "$job_id.2" {
			if {$debug} {send_user "\nmatch3\n"}
			incr matches
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: sstat not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$not_support != 0} {
		send_user "\nWARNING: can not test without accounting enabled\n"
		exit 0
	}
	if {$matches != 3} {
		send_user "\nFAILURE: sstat -$soption failed ($matches)\n"
		set exit_code 1
	}
	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} {send_user "\nmatch4\n"}
			incr matches
			exp_continue
		}
		-re "$job_id" {
			if {$debug} {send_user "\nmatch5\n"}
			incr matches
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: sstat not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$not_support != 0} {
		send_user "\nWARNING: can not test without accounting enabled\n"
		exit 0
	}
	if {$matches != 1} {
		send_user "\nFAILURE: sstat -$soption failed ($matches)\n"
		set exit_code 1
	}
	return $matches
}

	if { $soption == "-parsable" || $soption == "p" } {

	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} {send_user "\nmatch6\n"}
			incr matches
			exp_continue
		}
		-re "MaxRSSNode\\|MaxRSSTask\\|AveRSS\\|MaxPages\\|MaxPagesNode\\|MaxPagesTask\\|" {
			if {$debug} {send_user "\nmatch7\n"}
			incr matches
			exp_continue
		}
		-re "AvePages\\|MinCPU\\|MinCPUNode\\|MinCPUTask\\|AveCPU\\|NTasks\\|" {
			if {$debug} {send_user "\nmatch8\n"}
			incr matches
			exp_continue
		}
		-re "$job_id..\\|" {
			if {$debug} {send_user "\nmatch9\n"}
			incr matches
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: sstat not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$not_support != 0} {
		send_user "\nWARNING: can not test without accounting enabled\n"
		exit 0
	}
	if {$matches != 4} {
		send_user "\nFAILURE: sstat -$soption failed ($matches)\n"
		set exit_code 1
	}
	return $matches
}

	if { $soption == "-parsable2" || $soption == "P" } {

	spawn $sstat -$soption -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} {send_user "\nmatch10\n"}
			incr matches
			exp_continue
		}
		-re "MaxRSSNode\\|MaxRSSTask\\|AveRSS\\|MaxPages\\|MaxPagesNode\\|MaxPagesTask\\|" {
			if {$debug} {send_user "\nmatch11\n"}
			incr matches
			exp_continue
		}
		-re "AvePages\\|MinCPU\\|MinCPUNode\\|MinCPUTask\\|AveCPU\\|NTasks" {
			if {$debug} {send_user "\nmatch12\n"}
			incr matches
			exp_continue
		}
		-re "$job_id..\\|" {
			if {$debug} {send_user "\nmatch13\n"}
			incr matches
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: sstat not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$not_support != 0} {
		send_user "\nWARNING: can not test without accounting enabled\n"
		exit 0
	}
	if {$matches != 4} {
		send_user "\nFAILURE: sstat -$soption failed ($matches)\n"
		set exit_code 1
	}
		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} {send_user "\nmatch14\n"}
			incr matches
			exp_continue
		}
		-re "MaxRSSNode.MaxRSSTask.AveRSS.MaxPages.MaxPagesNode.MaxPagesTask" {
			if {$debug} {send_user "\nmatch15\n"}
			incr matches
			exp_continue
		}
		-re "AvePages.MinCPU.MinCPUNode.MinCPUTask.AveCPU.NTasks" {
			if {$debug} {send_user "\nmatch16\n"}
			incr matches
			exp_continue
		}
		-re "$job_id" {
			if {$debug} {send_user "\nmatch17\n"}
			incr matches
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: sstat not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$not_support != 0} {
		send_user "\nWARNING: can not test without accounting enabled\n"
		exit 0
	}
	if {$matches != 4} {
		send_user "\nFAILURE: sstat -$soption failed ($matches)\n"
		set exit_code 1
	}
		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 exit_code   0
	set matches     0
	set not_support 0
	send_user "sstat -$soption $vargs -p -j $job_id\n"

	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} {send_user "\nmatch18\n"}
			incr matches
			exp_continue
		}
		-re "JobID.MaxPages.MaxPagesNode.MaxPagesTask" {
			if {$debug} {send_user "\nmatch19\n"}
			incr matches
			exp_continue
		}
		-re "MaxRSS.MaxRSSNode.MaxRSSTask.MaxVMSize" {
			if {$debug} {send_user "\nmatch20\n"}
			incr matches
			exp_continue
		}
		-re "MaxVMSizeNode.MaxVMSizeTask.MinCPU.MinCPUNode" {
			if {$debug} {send_user "\nmatch21\n"}
			incr matches
			exp_continue
		}
		-re "MinCPUTask.NTasks.SystemCPU.TotalCPU" {
			if {$debug} {send_user "\nmatch22\n"}
			incr matches
			exp_continue
		}
		-re "$job_id" {
			if {$debug} {send_user "\nmatch23\n"}
			incr matches
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: sstat not responding\n"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$not_support != 0} {
		send_user "\nWARNING: can not test without accounting enabled\n"
		exit 0
	}
	if {$matches != 6} {
		send_user "\nFAILURE: sstat -$soption failed ($matches)\n"
		set exit_code 1
	}
		return $matches
	}
}
################################################################
sleep 1
set matches [sstat_job a $job_id3]
if {$matches != 3} {
	send_user "\nFAILURE: sstat -a failed ($matches)\n"
	set exit_code 1
}

set matches [sstat_job -allsteps $job_id3]
if {$matches != 3} {
	send_user "\nFAILURE: sstat --allsteps failed ($matches)\n"
	set exit_code 1
}

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

set matches [sstat_job -noheader $job_id3]
if {$matches != 1} {
	send_user "\nFAILURE: sstat --noheader failed ($matches)\n"
	set exit_code 1
}

set matches [sstat_job p $job_id3]
if {$matches != 4} {
	send_user "\nFAILURE: sstat -p failed ($matches)\n"
	set exit_code 1
}

set matches [sstat_job -parsable $job_id3]
if {$matches != 4} {
	send_user "\nFAILURE: sstat --parsable failed ($matches)\n"
	set exit_code 1
}

set matches [sstat_job P $job_id3]
if {$matches != 4} {
	send_user "\nFAILURE: sstat -P failed ($matches)\n"
	set exit_code 1
}

set matches [sstat_job -parsable2 $job_id3]
if {$matches != 4} {
	send_user "\nFAILURE: sstat --parsable2 failed ($matches)\n"
	set exit_code 1
}

set matches [sstat_job v $job_id3]
if {$matches != 4} {
	send_user "\nFAILURE: sstat -v failed ($matches)\n"
	set exit_code 1
}

set matches [sstat_job -verbose $job_id3]
if {$matches != 4} {
	send_user "\nFAILURE: sstat --verbose failed ($matches)\n"
	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,$sc,$tc $job_id3]
if {$matches != 6} {
	send_user "\nFAILURE: sstat -o failed ($matches)\n"
	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,$sc,$tc  $job_id3]
if {$matches != 6} {
	send_user "\nFAILURE: sstat --format failed ($matches)\n"
	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 {
		send_user "\nFAILURE: sacctmgr delete not responding\n"
		slow_kill $sadel_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$damatches != 1} {
	send_user "\nFAILURE: sacctmgr had a problem deleting account\n"
	set exit_code 1
}

if {$exit_code == 0} {
	exec $bin_rm -f $file_in1 $file_in2 $file_in3
	send_user "\nSUCCESS\n"
}
exit $exit_code
