#!/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 file_in1    "$test_dir/input"
set test_acct   "${test_name}_acct"
set timeout     $max_job_delay
set job_id     0
set user_name   [get_my_user_name]

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

proc cleanup {} {
	global job_id sacctmgr test_acct

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

#
# Use sacctmgr to add an account
#
set output [run_command -fail "$sacctmgr -i add account $test_acct"]
if [regexp "Nothing" $output] {
	fail "sacctmgr had a problem adding account"
}

#
# Add self to this new account
#
run_command -fail "$sacctmgr -i create user name=$user_name account=$test_acct"

#
# Build input script file1
#
make_bash_script $file_in1 "
  $srun --mem=20 --gres=craynetwork:0 --overlap sleep 60 &
  $srun --mem=20 --gres=craynetwork:0 --overlap sleep 60 &
  $srun --mem=20 --gres=craynetwork:0 --overlap sleep 60
"
#
# Submit a job using this account
# NOTE: --mem option here and in scripts above to permit running more than one step
#
set job_id [submit_job -fail "-N1 -t1 -v --mem=100 --gres=craynetwork:0 --account=$test_acct --output=/dev/null $file_in1"]
if {[get_job_param $job_id "Account"] != $test_acct} {
	fail "sbatch failed to use specified account"
}

#
# Wait for job to start running, then status it
#
wait_for_job -fail $job_id "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 eol
	set debug       0
	set matches     0
	set not_support 0
	log_debug "Testing sstat -$soption option"

	if { $soption == "-allsteps" || $soption == "a" } {
		set output [run_command_output -fail "$sstat -$soption -p -j $job_id"]
		incr matches [regexp -all "(?n)^$job_id\\.($number)" $output]
		incr matches [regexp -all "(?n)^$job_id\\.batch" $output]
		return $matches
	}

	if { $soption == "-noheader" || $soption == "n" } {
		set re_header "JobID|MaxVMSize|MaxVMSizeNode|MaxVMSizeTask|AveVMSize|MaxRSS|MaxRSSNode|MaxRSSTask|AveRSS|MaxPages|MaxPagesNode|MaxPagesTask|AvePages|MinCPU|MinCPUNode|MinCPUTask|AveCPU|NTasks"
		set output [run_command_output -fail "$sstat -$soption -p -j $job_id"]
		incr matches [regexp -all "(?n)^$job_id" $output]
		incr matches [regexp -all $re_header $output]
		return $matches
	}

	if { $soption == "-parsable" || $soption == "p" || $soption == "-parsable2" || $soption == "P" } {
		set output [run_command_output -fail "$sstat -$soption -j $job_id.$step_id"]
		incr matches [regexp -all "JobID\\|MaxVMSize\\|MaxVMSizeNode\\|MaxVMSizeTask\\|AveVMSize\\|MaxRSS\\|" $output]
		incr matches [regexp -all "MaxRSSNode\\|MaxRSSTask\\|AveRSS\\|MaxPages\\|MaxPagesNode\\|MaxPagesTask\\|" $output]
		incr matches [regexp -all "AvePages\\|MinCPU\\|MinCPUNode\\|MinCPUTask\\|AveCPU\\|NTasks\\|" $output]
		incr matches [regexp -all "\\|$eol" $output]
		incr matches [regexp -all "(?n)^$job_id\\.$step_id\\|" $output]
		return $matches
	}

	if { $soption == "-verbose" || $soption == "v" } {
		set output [run_command_output -fail "$sstat -$soption -p -j $job_id"]
		incr matches [regexp -all "JobID.MaxVMSize.MaxVMSizeNode.MaxVMSizeTask.AveVMSize.MaxRSS" $output]
		incr matches [regexp -all "MaxRSSNode.MaxRSSTask.AveRSS.MaxPages.MaxPagesNode.MaxPagesTask" $output]
		incr matches [regexp -all "AvePages.MinCPU.MinCPUNode.MinCPUTask.AveCPU.NTasks" $output]
		incr matches [regexp -all "(?n)^$job_id" $output]
		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 number
	set debug	0
	set matches     0
	set not_support 0
	log_debug "sstat -$soption $vargs -p -j $job_id"

	if { $soption == "o" || $soption == "-format" } {
		set output [run_command_output -fail "$sstat -$soption $vargs -p -j $job_id"]
		incr matches [regexp -all "AveCPU.AvePages.AveRSS.AveVMSize" $output]
		incr matches [regexp -all "JobID.MaxPages.MaxPagesNode.MaxPagesTask" $output]
		incr matches [regexp -all "MaxRSS.MaxRSSNode.MaxRSSTask.MaxVMSize" $output]
		incr matches [regexp -all "MaxVMSizeNode.MaxVMSizeTask.MinCPU.MinCPUNode" $output]
		incr matches [regexp -all "MinCPUTask.NTasks" $output]
		incr matches [regexp -all "(?n)^(?:\[^\\|\]*\\|){4}$job_id\\.($number)\\|" $output]
		return $matches
	}
}
################################################################
set matches [sstat_job a $job_id ""]
subtest {$matches == 4} "Verify sstat -a option" "$matches != 4"

set matches [sstat_job -allsteps $job_id ""]
subtest {$matches == 4} "Verify sstat --allsteps option" "$matches != 4"

set matches [sstat_job n $job_id ""]
subtest {$matches == 1} "Verify sstat -n option" "$matches != 1"

set matches [sstat_job -noheader $job_id ""]
subtest {$matches == 1} "Verify sstat --noheader option" "$matches != 1"

set matches [sstat_job p $job_id 0]
subtest {$matches == 6} "Verify sstat -p option" "$matches != 6"

set matches [sstat_job -parsable $job_id 0]
subtest {$matches == 6} "Verify sstat --parsable option" "$matches != 6"

set matches [sstat_job P $job_id 0]
subtest {$matches == 4} "Verify sstat -P option" "$matches != 4"

set matches [sstat_job -parsable2 $job_id 0]
subtest {$matches == 4} "Verify sstat --parsable2 option" "$matches != 4"

set matches [sstat_job v $job_id 0]
subtest {$matches == 4} "Verify sstat -v option" "$matches != 4"

set matches [sstat_job -verbose $job_id 0]
subtest {$matches == 4} "Verify sstat --verbose option" "$matches != 4"

set matches [sstat_vargs o $ac,$ap,$ar,$av,$ji,$mp,$mpn,$mpt,$mr,$mrn,$mrt,$mvs,$mvn,$mvt,$mc,$mn,$mt,$nt $job_id]
subtest {$matches == 6} "Verify sstat -o option" "$matches != 6"

set matches [sstat_vargs -format $ac,$ap,$ar,$av,$ji,$mp,$mpn,$mpt,$mr,$mrn,$mrt,$mvs,$mvn,$mvt,$mc,$mn,$mt,$nt $job_id]
subtest {$matches == 6} "Verify sstat --format option" "$matches != 6"
