#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Test srun stdout/err file name formatting (--output and --error
#          options with %j, %J, %n, %s and %t specifications).
############################################################################
# Copyright (C) 2002-2006 The Regents of the University of California.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Morris Jette <jette1@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_err_j       "$test_dir/j.%j.error"
set file_in          "$test_dir/input"
set file_out_perc    "$test_dir/%%.output"
set file_out_perc_result "$test_dir/%.output"
set file_out_J       "$test_dir/J.%J.output"
set file_out_n       "$test_dir/n.%n.output"
set file_out_s       "$test_dir/s.%s.output"
set file_out_t       "$test_dir/t.%t.output"
set file_out_u       "$test_dir/u.%u.output"
set file_err_u       "$test_dir/u.%u.error"
set user_name        ""
set job_id           0
set task_cnt         5

#
# Gets current user
#
set user_name [get_my_user_name]

proc cleanup {} {
	global job_id

	cancel_job $job_id
}

#
# Spawn a program that generates "task_id" (%t) in stdout file names
# and confirm they are created
#
set timeout $max_job_delay
spawn $srun --output=$file_out_t -N1 -n$task_cnt -O -v -t1 $bin_id
expect {
	-re "jobid ($number).*" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}
set file_cnt 0
for {set task_id 0} {$task_id < $task_cnt} {incr task_id} {
	set file_out_t_glob  "$test_dir/t.$task_id.output"
	if {[wait_for_file $file_out_t_glob] == 0} {
		incr file_cnt
	} else {
		log_error "File $file_out_t_glob was not created"
	}
	file delete $file_out_t_glob
}
subtest {$file_cnt == $task_cnt} "Spawn a program that generates task id (%t) in stdout file names and confirm they are created" "$task_cnt task output files were expected but $file_cnt were created"

#
# Spawn a shell via srun that includes "jobid" (%j) in stderr
# file name and confirm it is created
#
set job_id   0
set srun_exit 0
spawn $srun --error=$file_err_j --output=/dev/null -N1 -n$task_cnt -O -v -t1 $bin_sleep aaa
expect {
	-re "jobid ($number).*" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	-re "exit code $number" {
		log_debug "This error is expected, no worries"
		set srun_exit 1
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}
if {$srun_exit == 0} {
	fail "Exit code failed to indicate an error"
}

set file_err_j_glob  "$test_dir/j.$job_id.error"
subtest {[wait_for_file $file_err_j_glob] == 0} "Spawn a shell via srun that includes job name (%j) in stderr file name and confirm it is created"
file delete $file_err_j_glob

#
# Spawn a shell via srun that includes literal percent sign (%) in stdout
#  file name and confirm it is created
#
set job_id   0
spawn $srun --output=$file_out_perc -N1 -v -t1 $bin_printenv SLURMD_NODENAME
expect {
	-re "jobid ($number).*" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}

subtest {[wait_for_file $file_out_perc_result] == 0} "Spawn a shell via srun that includes literal percent sign (%%) in stdout file name and confirm it is created"
file delete $file_out_perc_result

#
# Spawn a shell via srun that includes "job_id.step_id" (%J) in stdout
#  file name and confirm it is created
#
set job_id   0
spawn $srun --output=$file_out_J -N1 -v -t1 $bin_printenv SLURMD_NODENAME
expect {
	-re "jobid ($number).*" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}

set task_id  0
set file_out_J_glob  "$test_dir/J.$job_id.$task_id.output"
subtest {[wait_for_file $file_out_J_glob] == 0} "Spawn a shell via srun that includes job step id (%J) in stdout file name and confirm it is created"
file delete $file_out_J_glob

#
# Spawn a shell via srun with the %u in stdout file name and confirm
# it is created
#
set job_id 0
spawn $srun --output=$file_out_u -v -t1 sleep 10
expect {
	-re "jobid ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun is not responding"
	}
	eof {
		wait
	}
}

if {$job_id == 0} {
	fail "srun did not submit job"
}

#
# Check that the correct files were created
#
set file_out_u_glob  "$test_dir/u.$user_name.output"
subtest {[wait_for_file $file_out_u_glob] == 0} "Spawn a shell via srun with %u in stdout file name and confirm it is created"
file delete $file_out_u_glob

#
# Spawn shell via srun with %u in the stderr file name and confirm
# it is created
#
set job_id 0
spawn $srun --error=$file_err_u --output=/dev/null -v -t1 sleep aaa
expect {
	-re "jobid ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	-re "error" {
		log_debug "Do not worry this error is expected"
		set srun_exit 1
		exp_continue
	}
	timeout {
		fail "srun is not responding"
	}
	eof {
		wait
	}
}

if {$job_id == 0} {
	fail "srun did not submit job"
}

#
# Check that the correct files where created
#
set file_err_u_glob  "$test_dir/u.$user_name.error"
subtest {[wait_for_file $file_err_u_glob] == 0} "Spawn shell via srun with %u in the stderr file name and confirm it is created"
file delete $file_err_u_glob

#
# Spawn a shell via srun that includes "node_id" (%n) in stdout
#  file name and confirm it is created
#
set node_id 0
set file_out_n_glob  "$test_dir/n.$node_id.output"
file delete $file_out_n_glob

set job_id   0
spawn $srun --output=$file_out_n -N1 -n$task_cnt -O -v -t1 $bin_printenv SLURMD_NODENAME
expect {
	-re "jobid ($number).*" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}
set file_cnt 0
for {set node_id 0} {$node_id < 2} {incr node_id} {
	set file_out_n_glob  "$test_dir/n.$node_id.output"

	if {($node_id == 0) && ([wait_for_file $file_out_n_glob] != 0)} {
		fail "Missing file ($file_out_n_glob)"
	} else {
		exec $bin_sleep 1
	}
	if [file exists $file_out_n_glob] {
		file delete $file_out_n_glob
		incr file_cnt
	}
}

subtest {$file_cnt == 1} "Confirm file format of %n in stdout file"

#
# Spawn a shell via srun that includes "step_id" (%s) in stdout
#  file name and confirm it is created. Use two step batch job.
#
make_bash_script $file_in "
  $srun -n4 -O --output=$file_out_s -t1 $bin_printenv SLURMD_NODENAME
  $srun -n4 -O --output=$file_out_s -t1 $bin_printenv SLURMD_NODENAME
"

for {set step_id 0} {$step_id < 4} {incr step_id} {
	set file_out_s_glob  "$test_dir/s.$step_id.output"
	file delete $file_out_s_glob
}

set node_cnt 1-4
set task_cnt 4

# Note: Task count is in the script for each srun
#	There is not oversubscribe option for sbatch.
set job_id   0
spawn $sbatch --output=/dev/null -N$node_cnt -t1 $file_in
expect {
	-re "Submitted batch job ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}

wait_for_job -fail $job_id DONE

set file_cnt 0
for {set step_id 0} {$step_id < 3} {incr step_id} {
	set file_out_s_glob  "$test_dir/s.$step_id.output"

	if {($step_id < 2) && ([wait_for_file $file_out_s_glob] != 0)} {
		fail "Missing file ($file_out_s_glob)"
	} else {
		exec $bin_sleep 1
	}
	if [file exists $file_out_s_glob] {
		file delete $file_out_s_glob
		incr file_cnt
	}
}

subtest {$file_cnt == 2} "Confirm format of %s in stdout file name ($file_cnt)"