#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Test of advanced reservation of licenses.
############################################################################
# Copyright (C) SchedMD LLC.
#
# 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 license_name        $test_name
set resv_name           "resv$test_id"
set user_name           ""
set job_id1             0
set job_id2             0

proc submit_job { license_count } {
	global bin_sleep license_name sbatch number
	set job_id 0
	spawn $sbatch -n1 -t1 -o /dev/null -L $license_name:$license_count --wrap "$bin_sleep 300"
	expect {
		-re "Submitted batch job ($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}
		timeout {
			fail "sbatch not responding"
		}
		eof {
			wait
		}
	}
	if { $job_id == 0 } {
		fail "Failed to submit job"
	}
	return $job_id
}

proc reason_is_licenses { job_id } {
	global squeue

	set reason_licenses 0
	spawn $squeue -j $job_id -o "%r"
	expect {
		-re "Licenses" {
			set reason_licenses 1
			exp_continue
		}
		timeout {
			fail "squeue not responding"
		}
		eof {
			wait
		}
	}
	return $reason_licenses
}

proc test_license_count {want_total want_used want_free } {
	global scontrol license_name number
	set license_free  0
	set license_total 0
	set license_used  0
	spawn $scontrol -o show license $license_name
	expect {
		-re "Total=($number)" {
			set license_total $expect_out(1,string)
			exp_continue
		}
		-re "Used=($number)" {
			set license_used $expect_out(1,string)
			exp_continue
		}
		-re "Free=($number)" {
			set license_free $expect_out(1,string)
			exp_continue
		}
		timeout {
			fail "scontrol not responding"
		}
		eof {
			wait
		}
	}
	if { $license_total != $want_total } {
		fail "Incorrect license_total count ($license_total != $want_total)"
	}
	if { $license_used != $want_used } {
		fail "Incorrect license_used count ($license_used != $want_used)"
	}
	if { $license_free != $want_free } {
		fail "Incorrect license_free count ($license_free != $want_free)"
	}
}

################# TEST STARTS HERE ######################
if {![is_super_user]} {
	skip "This test can't be run except as SlurmUser"
}

if {[check_config_select "linear"]} {
	skip "This test is incompatible with select/linear"
} elseif {[default_part_exclusive]} {
	skip "This test is incompatible with exclusive node allocations"
}

set user_name [get_my_user_name]

#
# Add licenses to system configuration, 8 licenses named "test3.15"
#
set got_config 0
set licenses   ""
log_user 0
spawn $scontrol show config
expect {
	-re "Licenses += .null" {
		exp_continue
	}
	-re "Licenses += ($re_word_str)" {
		set licenses $expect_out(1,string)
		exp_continue
	}
	-re "SLURM_CONF *= (/.*)/($re_word_str).*SLURM_VERSION" {
		set config_dir $expect_out(1,string)
		set got_config 1
		exp_continue
	}
	timeout {
		fail "scontrol is not responding"
	}
	eof {
		wait
	}
}
log_user 1
if {[string first ${license_name} ${licenses}] != -1} {
	fail "License ${license_name} already configured, likely vestigial from previous test, fix slurm.conf and test again"
}
if {$got_config == 0} {
	fail "Could not identify slurm.conf location"
}

set config_file $config_dir/slurm.conf
set cwd "[$bin_pwd]"

save_conf $config_file

proc cleanup {} {
	global job_id1 job_id2 config_file

	#
	# Cancel the jobs and
	# Restore the configuration
	#
	cancel_job [list $job_id1 $job_id2]

	restore_conf $config_file
	reconfigure
}

set sep ""
if {[string length ${licenses}] > 0} {
	set sep ","
}
exec $bin_grep -v Licenses $config_file > $cwd/slurm.conf.work
exec $bin_echo "Licenses=${licenses}${sep}${license_name}:8" >> $cwd/slurm.conf.work
exec $bin_cp $cwd/slurm.conf.work $config_file
exec $bin_rm -f $cwd/slurm.conf.work

reconfigure -fail

#
# Delete any vesgitial advanced reservation, then
# Create the advanced reservation with 6 of 8 licenses named "test3.15"
#
if { [delete_res $resv_name] } {
	log_info "Error is expected, no worries"
}

if { [create_res $resv_name "starttime=now duration=2 nodecnt=2 flags=license_only users=$user_name licenses=${license_name}:6"] != 0 } {
	fail "Error creating reservation"
}

#
# Submit batch job to claim 3 of 2 unreserved licenses of type "test3.15"
#
set job_id1 [submit_job 3]

#
# Check that job reason is "Licenses"
#
sleep 5
if { [reason_is_licenses $job_id1] == 0 } {
	fail "Failed to set proper job reason for job ($job_id1)"
}

#
# Drop reservation to only 2 of the "test3.15" licenses (leaving 6 licenses)
#
if { [update_res $resv_name "licenses=$license_name:2"] } {
	fail "Failed updating reservation licences"
}

#
# Check that job reason is no longer "Licenses"
#
sleep 5
if { [reason_is_licenses $job_id1] } {
	fail "Failed to set proper job reason for job ($job_id1)"
}

#
# Test scontrol show license output
#
test_license_count 8 3 5

#
# Submit batch job to claim 4 of 3 unreserved licenses of type "test3.15"
#
set job_id2 [submit_job 4]

#
# Check that job reason is "Licenses"
#
sleep 5
if { [reason_is_licenses $job_id2] == 0 } {
	fail "Failed to set proper job reason for job ($job_id2)"
}

#
# Drop reservation to only 1 of the "test3.15" licenses (leaving 7 licenses)
#
if { [update_res $resv_name "licenses=$license_name:1"] != 0 } {
	fail "Failed updating reservation licences"
}

#
# Check that job reason is no longer "Licenses"
#
sleep 5
if { [reason_is_licenses $job_id2] != 0 } {
	fail "Failed to set proper job reason for job ($job_id2)"
}

#
# Test scontrol show license output
#
test_license_count 8 7 1