#!/usr/bin/env expect
############################################################################
# Purpose: Test of setting and viewing SiteFactor
############################################################################
# Copyright (C) 2019 SchedMD LLC.
# Written by Brian Christiansen <brian@schedmd.com>
#
# 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 timeout	60

set jobid1 0
set jobid2 0

#
# 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 "PriorityType"] ne "priority/multifactor"} {
	skip "This test can't be run without a usable PriorityType"
}

#
# This test needs to be modified to use the core counts rather than CPU counts
#
set select_type_params [get_select_type_params ""]
if { [string first "CR_ONE_TASK_PER_CORE" $select_type_params] != -1 } {
	skip "This test can't be run SelectTypeParameters=CR_ONE_TASK_PER_CORE"
}

proc cleanup {} {
	global jobid1 jobid2

	cancel_job [list $jobid1 $jobid2]
}

proc sub_job { cpu_cnt state } {
	global number sbatch test_id bin_sleep

	set jobid 0
	spawn $sbatch -J$test_id -o/dev/null -e/dev/null -n $cpu_cnt --exclusive --wrap "$bin_sleep 300"
	expect {
		-re "Submitted batch job ($number)" {
			set jobid $expect_out(1,string)
			exp_continue
		}
		timeout {
			fail "sbatch not responding"
		}
		eof {
			wait
		}
	}

	if {$jobid == 0} {
		fail "Did not get sbatch jobid"
	}

	wait_for_job -fail $jobid $state

	return $jobid
}


proc sprio_args { args regex match_cnt } {
	global sprio

	if [wait_for_command "$sprio $args" "\[regexp -all {$regex} \[dict get \$result output\]\] == $match_cnt"] {
		fail "sprio did not match"
	}
}

proc update_job { jobid prio {error ""} } {
	global scontrol

	set matches 0

	spawn $scontrol update jobid=$jobid sitefactor=$prio
	set tmp_re ".*"
	if {$error != ""} {
		set tmp_re $error
	}
	expect {
		-re $tmp_re {
			incr matches
		}
		timeout {
			fail "scontrol not responding"
		}
		eof {
			wait
		}
	}

	if {($error != "") && !$matches} {
		fail "Didn't get expected error '$error'"
	}
}

################################################################
set cpu_cnt [get_total_cpus]
set jobid1 [sub_job $cpu_cnt RUNNING]
set jobid2 [sub_job $cpu_cnt PENDING]

if {[info exists env(SPRIO_FORMAT)]} {
    unset env(SPRIO_FORMAT)
}

sprio_args "-j $jobid2 -o \"%.15i %S\"" "$jobid2 \[-\]*$number" 1
update_job $jobid2 100
sprio_args "-j $jobid2 -o \"%.15i %S\"" "$jobid2 100" 1
update_job $jobid2 -100
sprio_args "-j $jobid2 -o \"%.15i %S\"" "$jobid2 -100" 1
update_job $jobid2 0
sprio_args "-j $jobid2 -o \"%.15i %S\"" "$jobid2 0" 1

update_job $jobid2 2147483645
sprio_args "-j $jobid2 -o \"%.15i %S\"" "$jobid2 2147483645" 1
update_job $jobid2 -2147483645
sprio_args "-j $jobid2 -o \"%.15i %S\"" "$jobid2 -2147483645" 1

update_job $jobid2 2147483646 "scontrol: error: SiteFactor value out of range"
update_job $jobid2 -2147483646 "scontrol: error: SiteFactor value out of range"
