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

proc cleanup {} {
	global sacctmgr test_acct

	run_command "$sacctmgr -i delete account $test_acct"
}

set accounting_storage_type [get_config_param "AccountingStorageType"]

set cluster [get_config_param "ClusterName"]
#
# Identify the user and his group
#
set user_name [get_my_user_name]
set user_gid [get_my_gid]

#
# Use sacctmgr to add an account
#
set aamatches 0
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 {
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}
if {$aamatches != 1} {
	fail "sacctmgr had a problem adding account"
}

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


#
# Spawn a job via srun using this account
#
set output [run_command_output -fail "$srun -N1 -v --account=$test_acct $bin_id"]
if {![regexp "launching StepId=($number)\.$re_word_str" $output - job_id]} {
	fail "Did not get srun job id"
}
subtest {[get_job_param $job_id "Account"] eq $test_acct} "Verify srun job is using the specified account"

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 eq "-brief" || $soption eq "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 eq "-long" || $soption eq "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 eq "-noheader" || $soption eq "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 eq "-parsable" || $soption eq "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 eq "-parsable2" || $soption eq "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 eq "-verbose" || $soption eq "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 eq "g" || $soption eq "-gid" || $soption eq "-group" || $soption eq "u" || $soption eq "-uid" || $soption eq "-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
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 {
		fail "sacctmgr add not responding"
	}
	eof {
		wait
	}
}
if {$aamatches != 1} {
	fail "sacctmgr had a problem adding account"
}

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

#
# Spawn a job via srun using this account
#
set output [run_command_output -fail "$srun -N1 -v --account=$test_acct $bin_id"]
if {![regexp "launching StepId=($number)\\.$re_word_str" $output - job_id]} {
	fail "Did not get srun job id"
}
subtest {[get_job_param $job_id "Account"] eq $test_acct} "Verify srun job is using the specified account"

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]
subtest {$matches == 2} "Verify sacct -b option" "sacct -b failed ($matches != 2)"

set matches [sacct_job -brief $job_id]
subtest {$matches == 2} "Verify sacct --brief option" "sacct --brief failed ($matches != 2)"

set matches [sacct_vargs g $user_gid $job_id]
subtest {$matches == 3} "Verify sacct -g option" "sacct -g failed ($matches != 3)"

set matches [sacct_vargs -gid $user_gid $job_id]
subtest {$matches == 3} "Verify sacct --gid option" "sacct --gid failed ($matches != 3)"

set matches [sacct_vargs -group $user_gid $job_id]
subtest {$matches == 3} "Verify sacct --group option" "sacct --group failed ($matches != 3)"

set matches [sacct_job l $job_id]
subtest {$matches == 8} "Verify sacct -l option" "sacct -l failed ($matches != 8)"

set matches [sacct_job -long $job_id]
subtest {$matches == 8} "Verify sacct --long option" "sacct --long failed ($matches != 8)"

set matches [sacct_job n $job_id]
subtest {$matches == 1} "Verify sacct -n option" "sacct -n failed ($matches != 1)"

set matches [sacct_job -noheader $job_id]
subtest {$matches == 1} "Verify sacct --noheader option" "sacct --noheader failed ($matches != 1)"

set matches [sacct_job p $job_id]
subtest {$matches == 2} "Verify sacct -p option" "sacct -p failed ($matches != 2)"

set matches [sacct_job -parsable $job_id]
subtest {$matches == 2} Verify sacct --parsable option" "sacct --parsable failed ($matches != 2)"

set matches [sacct_job P $job_id]
subtest {$matches == 2} "Verify sacct -P option" "sacct -P failed ($matches != 2)"

set matches [sacct_job -parsable2 $job_id]
subtest {$matches == 2} "Verify sacct --parsable2 option" "sacct --parsable2 failed ($matches != 2)"

set matches [sacct_vargs u $user_name $job_id]
subtest {$matches == 3} "Verify sacct -u option" "sacct -g failed ($matches != 3)"

set matches [sacct_vargs -uid $user_name $job_id]
subtest {$matches == 3} "Verify sacct --uid option" "sacct --gid failed ($matches != 3)"

set matches [sacct_vargs -user $user_name $job_id]
subtest {$matches == 3} "Verify sacct --user option" "sacct --group failed ($matches != 3)"

set matches [sacct_job v $job_id]
subtest {$matches == 4} "Verify sacct -v option" "sacct -v failed ($matches != 4)"

set matches [sacct_job -verbose $job_id]
subtest {$matches == 4} "Verify sacct --verbose option" "sacct --verbose failed ($matches != 4)"