#!/usr/bin/env expect
############################################################################
# Purpose: Test of sacct functionality for array jobs.
#          sacct options j
############################################################################
# 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 array1_job_id   0
set array2_job_id   0
set job_array_size  6
set task_num        2
set timeout $max_job_delay

#
# 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"
}

proc cleanup {} {
	global array1_job_id array2_job_id

	cancel_job [list $array1_job_id $array2_job_id]
}

proc run_sacct { jobs exp_match_1 exp_match_2 } {
	global sacct expected_steps array1_job_id array2_job_id

	# Now test sacct
	set output [run_command_output -fail "$sacct -j $jobs --starttime=00:00 --format JobID --noheader"]
	subtest [regexp $array1_job_id $output] "Verify job record for array job 1 was found in the accounting database"
	subtest [regexp $array2_job_id $output] "Verify job record for array job 2 was found in the accounting database"
}

set array1_job_id [submit_job -fail "-N1 --array=0-[expr $job_array_size - 1] --output=/dev/null -t1 --wrap=exit"]
set array2_job_id [submit_job -fail "-N1 --array=0-[expr $job_array_size - 1] --output=/dev/null -t1 --wrap=exit"]

# If PrologFlags=contain is in slurm.conf, then an "extern" step will be
# launched on each node, so we need to check for 3 steps per job instead of 2.
if {[param_contains [get_config_param "PrologFlags"] "Contain"]} {
	set expected_steps 3
} else {
	set expected_steps 2
}

# Wait for the accounting records for all job array elements to reach slurmdbd
wait_for_command -fail "$sacct -j $array1_job_id --allocation --format state --parsable2 --noheader" "\[regexp -all {COMPLETED} \[dict get \$result output\]\] == $job_array_size"
wait_for_command -fail "$sacct -j $array2_job_id --allocation --format state --parsable2 --noheader" "\[regexp -all {COMPLETED} \[dict get \$result output\]\] == $job_array_size"

# Now test sacct one task from each array
run_sacct "$array1_job_id\_1,$array2_job_id\_3" $expected_steps $expected_steps

# Now test sacct first full array and one task from the second array
run_sacct "$array1_job_id,$array2_job_id\_3" [expr $expected_steps * $job_array_size] $expected_steps

# Now test sacct one task from the first array and the second full
run_sacct "$array1_job_id\_2,$array2_job_id" $expected_steps [expr $expected_steps * $job_array_size]