#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm sacct functionality
#          sacct options b, g, j, l, n, p, u, v.
############################################################################
# Copyright (C) 2008 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 test_acct   "test$test_id\_acct"
set timeout 60

#
# 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 {[string compare [get_admin_level] "Administrator"]} {
	skip "This test can't be run without being an Accounting administrator. Use: sacctmgr mod user \$USER set admin=admin"
}

set accounting_storage_type [get_config_param "AccountingStorageType"]

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

set user_gid [get_my_gid]

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 cluster=$cluster]
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
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}
if {$aamatches != 1} {
	fail "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 cluster=$cluster]
expect {
	 timeout {
		slow_kill $sadd_pid
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}


#
# Spawn a job via srun using this account
#
set job_id 0
spawn $srun -N1 -v --account=$test_acct $bin_id
expect {
	-re "launching StepId=($number)\\.$re_word_str" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	log_error "did not get srun job_id3"
	set exit_code 1
} else {
	set matches 0
	spawn $scontrol show job $job_id
	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
	}
}

if {$accounting_storage_type eq "accounting_storage/slurmdbd"} {
	sleep 12
}

################################################################
#
# Proc: sacct_job
#
# Purpose:  Pass sacct options and test
#
# Returns: Number of matches.
#
# Input: Switch options not requiring arguments
#
################################################################

proc sacct_job { soption job_id} {
	global sacct accounting_storage_type

	set debug       0
	set matches     0
	set not_support 0
	set total_opts "$sacct -$soption -X -p --starttime=00:00 -j $job_id"

	eval spawn $total_opts

	if { $soption == "-brief" || $soption == "b" } {
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID.State.ExitCode" {
				if {$debug} {log_debug "Match1"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				if {$debug} {log_debug "Match2"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sacct not responding"
			}
			eof {
				wait
			}
		}

		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-long" || $soption == "l" } {
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID.JobIDRaw.JobName.Partition.MaxVMSize" {
				if {$debug} {log_debug "Match3"}
				incr matches
				exp_continue
			}
			-re "MaxVMSizeNode.MaxVMSizeTask.AveVMSize.MaxRSS" {
				if {$debug} {log_debug "Match4"}
				incr matches
				exp_continue
			}
			-re "MaxRSSNode.MaxRSSTask.AveRSS.MaxPages" {
				if {$debug} {log_debug "Match5"}
				incr matches
				exp_continue
			}
			-re "MaxPagesNode.MaxPagesTask.AvePages.MinCPU" {
				if {$debug} {log_debug "Match6"}
				incr matches
				exp_continue
			}
			-re "MinCPUNode.MinCPUTask.AveCPU.NTasks" {
				if {$debug} {log_debug "Match7"}
				incr matches
				exp_continue
			}
			-re "AllocCPUS.Elapsed.State.ExitCode" {
				if {$debug} {log_debug "Match8"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				if {$debug} {log_debug "Match9"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sacct not responding"
			}
			eof {
				wait
			}
		}
		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-noheader" || $soption == "n" } {
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "AllocCPUS|Account|AssocID|AveCPU|AvePages|AveRSS|AveVSize|BlockID	\
				|Cluster|CPUTime|CPUTimeRAW|Elapsed	\
				|Eligible|End|ExitCode|GID	\
				|Group|JobID|JobName|NodeList	\
				|MaxPages|MaxPagesNode|MaxPagesTask|MaxRSS	|
				|MaxRSSNode|MaxRSSTask|MaxVSize|MaxVSizeNode	|
				|MaxVSizeTask|MinCPU|MinCPUNode|MinCPUTask	|
				|NCPUS|NNodes|NTasks|Priority	|
				|Partition|QOS|QOSRAW|ReqCPUS	|
				|Reserved|ResvCPU|ResvCPURAW|Start	|
				|State|Submit|Suspended|SystemCPU	|
				|Timelimit|TotalCPU|UID|User	|
				|UserCPU|WCKey|WCKeyID" {
				if {$debug} {log_debug "Match10"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				if {$debug} {log_debug "Match11"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sacct not responding"
			}
			eof {
				wait
			}
		}

		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-parsable" || $soption == "p" } {
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID\\|JobName\\|Partition\\|Account\\|AllocCPUS\\|State\\|ExitCode\\|" {
				if {$debug} {log_debug "Match12"}
				incr matches
				exp_continue
			}
			-re "$job_id\\|" {
				if {$debug} {log_debug "Match13"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sacct not responding"
			}
			eof {
				wait
			}
		}

		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-parsable2" || $soption == "P" } {
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID\\|JobName\\|Partition\\|Account\\|AllocCPUS\\|State\\|ExitCode *" {
				if {$debug} {log_debug "Match14"}
				incr matches
				exp_continue
			}
			-re "$job_id\\|" {
				if {$debug} {log_debug "Match15"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sacct not responding"
			}
			eof {
				wait
			}
		}

		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}

	if { $soption == "-verbose" || $soption == "v" } {
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "sacct: accounting_storage/slurmdbd: init: Accounting storage SLURMDBD plugin loaded" {
				if {$debug} {log_debug "Match16"}
				incr matches
				exp_continue
			}
			-re "JobID.JobName.Partition" {
				if {$debug} {log_debug "Match17"}
				incr matches
				exp_continue
			}
			-re "Account.AllocCPUS.State.ExitCode" {
				if {$debug} {log_debug "Match18"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				if {$debug} {log_debug "Match19"}
				incr matches
				exp_continue
			}
			timeout {
				fail "sacct not responding"
			}
			eof {
				wait
			}
		}

		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}
}

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

proc sacct_vargs { soption vargs job_id} {
	global sacct

	set debug       0
	set matches     0
	set not_support 0
	set total_opts "$sacct -$soption $vargs -X -p -j $job_id"

	eval spawn $total_opts

	if { $soption == "g" || $soption == "-gid" || $soption == "-group" || $soption == "u" || $soption == "-uid" || $soption == "-user"} {
		expect {
			-re "Slurm accounting storage is disabled" {
				set not_support 1
				exp_continue
			}
			-re "JobID.JobName.Partition" {
				if {$debug} {log_debug "Match20"}
				incr matches
				exp_continue
			}
			-re "Account.AllocCPUS.State.ExitCode" {
				if {$debug} {log_debug "Match21"}
				incr matches
				exp_continue
			}
			-re "$job_id" {
				incr matches
				if {$debug} {log_debug "Match22"}
				exp_continue
			}
			timeout {
				fail "sacct not responding"
			}
			eof {
				wait
			}
		}

		if {$not_support != 0} {
			skip "Can not test without accounting enabled"
		}
		return $matches
	}
}
################################################################

#
# Use sacctmgr to add an account
#
set aamatches 0
set sadd_pid [spawn $sacctmgr -i add account $test_acct cluster=$cluster]
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
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}
if {$aamatches != 1} {
	fail "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 cluster=$cluster]
expect {
	 timeout {
		slow_kill $sadd_pid
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}

#
# Spawn a job via srun using this account
#
set job_id 0
spawn $srun -N1 -v --account=$test_acct $bin_id
expect {
	-re "launching StepId=($number)\\.$re_word_str" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	log_error "Did not get srun job_id3"
	set exit_code 1
} else {
	set matches 0
	spawn $scontrol show job $job_id
	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
	}

	if {$accounting_storage_type eq "accounting_storage/slurmdbd"} {
		sleep 12
	}

#
#	Wait for accounting data to be propagated to slurmdbd
#
	sleep 10

	set matches [sacct_job b $job_id]
	if {$matches != 2} {
		log_error "sacct -b failed ($matches != 2)"
		set exit_code 1
	}

	set matches [sacct_job -brief $job_id]
	if {$matches != 2} {
		log_error "sacct --brief failed ($matches != 2)"
		set exit_code 1
	}

	set matches [sacct_vargs g $user_gid $job_id]
	if {$matches != 3} {
		log_error "sacct -g failed ($matches != 3)"
		set exit_code 1
	}

	set matches [sacct_vargs -gid $user_gid $job_id]
	if {$matches != 3} {
		log_error "sacct --gid failed ($matches != 3)"
		set exit_code 1
	}

	set matches [sacct_vargs -group $user_gid $job_id]
	if {$matches != 3} {
		log_error "sacct --group failed ($matches != 3)"
		set exit_code 1
	}

	set matches [sacct_job l $job_id]
	if {$matches != 8} {
		log_error "sacct -l failed ($matches != 8)"
		set exit_code 1
	}

	set matches [sacct_job -long $job_id]
	if {$matches != 8} {
		log_error "sacct --long failed ($matches != 8)"
		set exit_code 1
	}

	set matches [sacct_job n $job_id]
	if {$matches != 1} {
		log_error "sacct -n failed ($matches != 1)"
		set exit_code 1
	}

	set matches [sacct_job -noheader $job_id]
	if {$matches != 1} {
		log_error "sacct -n failed ($matches != 1)"
		set exit_code 1
	}

	set matches [sacct_job p $job_id]
	if {$matches != 2} {
		log_error "sacct -p failed ($matches != 2)"
		set exit_code 1
	}

	set matches [sacct_job -parsable $job_id]
	if {$matches != 2} {
		log_error "sacct --parsable failed ($matches != 2)"
		set exit_code 1
	}

	set matches [sacct_job P $job_id]
	if {$matches != 2} {
		log_error "sacct -P failed ($matches != 2)"
		set exit_code 1
	}

	set matches [sacct_job -parsable2 $job_id]
	if {$matches != 2} {
		log_error "sacct --parsable2 failed ($matches != 2)"
		set exit_code 1
	}

	set matches [sacct_vargs u $user_name $job_id]
	if {$matches != 3} {
		log_error "sacct -g failed ($matches != 3)"
		set exit_code 1
	}

	set matches [sacct_vargs -uid $user_name $job_id]
	if {$matches != 3} {
		log_error "sacct --gid failed ($matches != 3)"
		set exit_code 1
	}

	set matches [sacct_vargs -user $user_name $job_id]
	if {$matches != 3} {
		log_error "sacct --group failed ($matches != 3)"
		set exit_code 1
	}

	set matches [sacct_job v $job_id]
	if {$matches != 4} {
		log_error "sacct -v failed ($matches != 4)"
		set exit_code 1
	}

	set matches [sacct_job -verbose $job_id]
	if {$matches != 4} {
		log_error "sacct --verbose failed ($matches != 4)"
		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 a problem deleting account"
}

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