#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          sacctmgr verify wckey
############################################################################
# Copyright (C) 2008-2010 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 cwd          "[$bin_pwd]"
set file_in     "test.$test_id.input"
set ta		"test$test_id-account.1"
set wk          "test$test_id-wckey"
set wk2         "test$test_id-wckey2"
set wk3         "test$test_id-wckey3"
set timeout 60

set config_path [get_conf_path]
fail_on_error "Error getting config path"
set config_file $config_path/slurm.conf

#
# Check accounting config and bail if not found.
#
if { [test_account_storage] == 0 } {
	skip "This test can't be run without a usable AccountStorageType"
}

if { [string compare [check_accounting_admin_level] "Administrator"] } {
	skip "This test can't be run without being an Accounting administrator. Use: sacctmgr mod user \$USER set admin=admin."
}

set using_slurmdbd [test_using_slurmdbd]

set cn [get_cluster_name]

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

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 $ta cluster=$cn]
expect {
	-re "Adding Account" {
		incr aamatches
		exp_continue
	}
	-re "Nothing new added" {
		log_warn "Vestigial account $ta found"
		incr aamatches
		exp_continue
	}
	timeout {
		log_error "sacctmgr add not responding"
		slow_kill $sadd_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$aamatches != 1} {
	fail "sacctmgr had a problem adding account."
}

#
# Add self to this new account and wckey
#
set sadd_pid [spawn $sacctmgr -i create user name=$user_name account=$ta cluster=$cn wckeys=$wk]
expect {
	 timeout {
		log_error "sacctmgr add not responding"
		slow_kill $sadd_pid
		set exit_code 1
	}
	eof {
		wait
	}
}

#
# Enable TrackWCKey
#
set reconfigured false
if {![test_track_wckey_slurmctld]} {
	log_info "Enabling TrackWCKey..."

	save_conf $config_file

	exec $bin_sed -i /^\[\t\s\]*TrackWCKey\[\t\s\]*=/Id $config_file
	exec $bin_echo TrackWCKey=Yes >> $config_file
	reconfigure
	set reconfigured true
}

#
# Spawn a job via salloc using this wckey
#
set job_id     0
set matches    0
set timeout    $max_job_delay
spawn $salloc -N1 --account=$ta --wckey=$wk
expect {
	-re "Granted job allocation ($number)" {
		set job_id $expect_out(1,string)
		send "$scontrol show job $job_id\r"
		send "exit\r"
		exp_continue
	}
	-re "WCKey=$wk" {
		incr matches
		exp_continue
	}
	timeout {
		log_error "salloc not responding"
		set exit_code 1
	}
	eof {
		wait
	}
}

if {$using_slurmdbd} {
	sleep 5
}

spawn $sacct -p -j $job_id --format=wckey
	expect {
		-re "$wk" {
			incr matches
			exp_continue
		}
		timeout {
			log_error "scontrol not responding"
			set exit_code 1
		}
		eof {
			wait
		}
	}
if {$job_id == 0} {
	log_error "salloc failed to initiate job"
	set exit_code 1
} elseif {$matches != 2} {
	log_error "salloc failed to use desired wckey"
	set exit_code 1
}

#
# Spawn a job via sbatch using this wckey
#
make_bash_script $file_in "$bin_id"
set job_id 0
spawn $sbatch -N1 --account=$ta --wckey=$wk --output=none $file_in
expect {
	-re "Submitted batch job ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		log_error "sbatch not responding"
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	log_error "did not get sbatch job_id"
	set exit_code 1
} else {
	set matches 0
	spawn $scontrol show job $job_id
	expect {
		-re "WCKey=$wk" {
			incr matches
			exp_continue
		}
		timeout {
			log_error "scontrol not responding"
			set exit_code 1
		}
		eof {
			wait
		}
	}
	if {$using_slurmdbd} {
		sleep 5
	}

	spawn $sacct -p -j $job_id --format=wckey
	expect {
		-re "$wk" {
			incr matches
			exp_continue
		}
		timeout {
			log_error "scontrol not responding"
			set exit_code 1
		}
		eof {
			wait
		}
	}
	if {$matches != 2} {
		log_error "sbatch failed to use specified wckey"
		set exit_code 1
	}
}

#
# Spawn a job via srun using this wckey
#
set job_id 0
spawn $srun -N1 -v --account=$ta --wckey=$wk $bin_id
expect {
	-re "launching ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		log_error "srun not responding"
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	log_error "did not get srun job_id"
	set exit_code 1
} else {
	set matches 0
	spawn $scontrol show job $job_id
	expect {
		 -re "WCKey=$wk" {
			incr matches
			exp_continue
		}
		timeout {
			log_error "scontrol not responding"
			set exit_code 1
		}
		eof {
			wait
		}
	}

	if {$using_slurmdbd} {
		sleep 5
	}

	spawn $sacct -p -j $job_id --format=wckey
	expect {
		-re "$wk" {
			incr matches
			exp_continue
		}
		timeout {
			log_error "scontrol not responding"
			set exit_code 1
		}
		eof {
			wait
		}
	}
	if {$matches != 2} {
		log_error "srun failed to use specified wckey"
		set exit_code 1
	}
}

#
# Use sacctmgr to delete the test account
#
set damatches 0
set sadel_pid [spawn $sacctmgr -i delete account $ta]
expect {
	-re "Deleting account" {
		incr damatches
		exp_continue
	}
	timeout {
		log_error "sacctmgr delete not responding"
		slow_kill $sadel_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$damatches != 1} {
	log_error "sacctmgr had a problem deleting account"
	set exit_code 1
}

#
# Add more test WCKeys the user
#
set awmatches 0
set saadd_pid [spawn $sacctmgr -i add user $user_name cluster=$cn wckey=$wk,$wk2,$wk3]

expect {
	-re "W ="	{
		incr awmatches
		exp_continue
	}
	-re "Nothing new added" {
		log_error "WCKeys used in test already existed"
		set exit_code 1
	}
	timeout {
		log_error "sacctmgr add not responding"
		slow_kill $saadd_pid
		set exit_code 1
	}
	eof {
		wait
	}
}

# $wk should be there before, so only $wk2 and $wk3 was added
if {$awmatches!=2} {
	log_error "sacctmgr had problems adding two WCKey at once"
	set exit_code 1
}

#
# Verify that we have 3 WCKeys in DB
#
set swmatches 0
set salist_pid [spawn $sacctmgr list -p wckey user=$user_name]
expect {
	-ex "\n$wk\|$cn\|$user_name\|"  {
		incr swmatches
		exp_continue
	}
	-ex "\n$wk2\|$cn\|$user_name\|" {
		incr swmatches
		exp_continue
	}
	-ex "\n$wk3\|$cn\|$user_name\|" {
		incr swmatches
		exp_continue
	}
	timeout {
		log_error "sacctmgr list not responding"
		slow_kill $salist_pid
		set exit_code 1
	}
	eof {
		wait
	}
}

if {$swmatches!=3} {
	log_error "We should have 3 test WCKeys configured we have $swmatches"
	set exit_code 1
}

#
# Use sacctmgr to delete test WCKeys
#
set dwmatches 0
set sadel_pid [spawn $sacctmgr -i delete user $user_name wckey=$wk,$wk2,$wk3]
expect {
	-re "Deleteing WCKeys" {
		incr dwmatches
		exp_continue
	}
	timeout {
		log_error "sacctmgr delete not responding"
		slow_kill $sadel_pid
		set exit_code 1
	}
	eof {
		wait
	}
}

#
# Use sacctmgr to verify if WCKeys were removed
#
set salist_pid [spawn $sacctmgr -p list wckey where user=$user_name]
expect {
	-ex "\n$wk\|$cn\|$user_name\|" {
		set exit_code 1
		log_error "WCKey = $wk should be deleted, but was not"
	}
	-ex "\n$wk2\|$cn\|$user_name\|" {
		set exit_code 1
		log_error "WCKey = $wk2 should be deleted, but was not"
	}
	-ex "\n$wk3\|$cn\|$user_name\|" {
		set exit_code 1
		log_error "WCKey = $wk3 should be deleted, but was not"
	}
	timeout {
		log_error "sacctmgr list not responding"
		slow_kill $salist_pid
		set exit_code 1
	}
	eof {
		wait
	}
}

if {$reconfigured} {
	restore_conf $config_file
	reconfigure
}

if {$exit_code == 0} {
	exec $bin_rm -f $file_in
} else {
	fail "Test failed due to previous errors (\$exit_code = $exit_code)"
}