#!/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 exit_code        0
set file_err_j       "test$test_id.j.%j.error"
set file_in          "test$test_id.input"
set file_out_perc    "test$test_id.%%.output"
set file_out_perc_result "test$test_id.%.output"
set file_out_J       "test$test_id.J.%J.output"
set file_out_n       "test$test_id.n.%n.output"
set file_out_s       "test$test_id.s.%s.output"
set file_out_t       "test$test_id.t.%t.output"
set file_out_u       "test$test_id.u.%u.output"
set file_err_u       "test$test_id.u.%u.error"
set user_name        ""
set job_id           0
set task_cnt         5

#
# Gets current user
#

set user_name [get_my_user_name]

#
# Spawn a program that generates "task_id" (%t) in stdout file names
# and confirm they are created
#
for {set task_id 0} {$task_id < $task_cnt} {incr task_id} {
	set file_out_t_glob  "test$test_id.t.$task_id.output"
	exec $bin_rm -f $file_out_t_glob
}
set timeout $max_job_delay
set srun_pid [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 {
		slow_kill $srun_pid
		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$test_id.t.$task_id.output"
	if {[wait_for_file $file_out_t_glob] != 0} {
		set exit_code 1
	} else {
		exec $bin_rm -f $file_out_t_glob
		incr file_cnt
	}
}
if {$file_cnt != $task_cnt} {
	log_error "File format of %t in stdout failed"
	set exit_code 1
}

#
# 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
set srun_pid [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 {
		slow_kill $srun_pid
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}
if {$srun_exit == 0} {
	log_error "Exit code failed to indicate an error"
	set exit_code 1
}

set file_err_j_glob  "test$test_id.j.$job_id.error"
if {[wait_for_file $file_err_j_glob] == 0} {
	exec $bin_rm -f $file_err_j_glob
} else {
	log_error "File format of %j in stderr failed"
	set exit_code 1
}

#
# Spawn a shell via srun that includes literal percent sign (%) in stdout
#  file name and confirm it is created
#
set job_id   0
set srun_pid [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 {
		slow_kill $srun_pid
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}
if {[wait_for_file $file_out_perc_result] != 0} {
	log_error "File format of %% in stdout failed"
	set exit_code 1
} else {
	exec $bin_rm -f $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
set srun_pid [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 {
		slow_kill $srun_pid
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
}
set task_id  0
set file_out_J_glob  "test$test_id.J.$job_id.$task_id.output"
if {[wait_for_file $file_out_J_glob] != 0} {
	log_error "File format of %J in stdout failed"
	set exit_code 1
} else {
	exec $bin_rm -f $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
set srun_pid [spawn $srun --output=$file_out_u -v -t1 sleep 10]
expect {
	-re "jobid ($number)" {
		set job_id $expect_out(1,string)
		exp_continue
	}
	timeout {
		slow_kill $srun_pid
		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$test_id.u.$user_name.output"

if {[wait_for_file $file_out_u_glob]==0} {
	exec $bin_rm -f $file_out_u_glob
} else {
	log_error "File format is not created correctly"
	set exit_code 1
}

#
# Spawn shell via srun with %u in the stderr file name and confirm
# it is created
#
set job_id 0
set srun_pid [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 {
		slow_kill $srun_pid
		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$test_id.u.$user_name.error"

if {[wait_for_file $file_err_u_glob]==0} {
	exec $bin_rm -f $file_err_u_glob
} else {
	log_error "File format is not created correctly"
	set exit_code 1
}

#
# 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$test_id.n.$node_id.output"
exec $bin_rm -f $file_out_n_glob

set job_id   0
set srun_pid [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 {
		slow_kill $srun_pid
		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$test_id.n.$node_id.output"

	if {($node_id == 0) && ([wait_for_file $file_out_n_glob] != 0)} {
		log_error "Missing file $file_out_n_glob"
		set exit_code 1
	} else {
		exec $bin_sleep 1
	}
	if [file exists $file_out_n_glob] {
		exec $bin_rm -f $file_out_n_glob
		incr file_cnt
	}
}

if {$file_cnt != 1} {
	log_error "File format of %n in stdout failed"
	set exit_code 1
}

#
# Spawn a shell via srun that includes "step_id" (%s) in stdout
#  file name and confirm it is created. Use two step batch job.
#
exec $bin_rm -f $file_in
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$test_id.s.$step_id.output"
	exec $bin_rm -f $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
set srun_pid [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 {
		slow_kill $srun_pid
		fail "srun not responding"
	}
	eof {
		wait
	}
}
if {$job_id == 0} {
	fail "Job initiation failed"
} else {
	if {[wait_for_job $job_id DONE] != 0} {
		log_error "Error completing job $job_id"
		cancel_job $job_id
		set exit_code 1
	}
}
exec $bin_rm -f $file_in

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

	if {($step_id < 2) && ([wait_for_file $file_out_s_glob] != 0)} {
		log_error "Missing file $file_out_s_glob"
		set exit_code 1
	} else {
		exec $bin_sleep 1
	}
	if [file exists $file_out_s_glob] {
		exec $bin_rm -f $file_out_s_glob
		incr file_cnt
	}
}

if {$file_cnt != 2} {
	log_error "File format of %s in stdout failed($file_cnt)"
	set exit_code 1
}

#
# Post-processing
#
if {$exit_code != 0} {
	fail "Test failed due to previous errors (\$exit_code = $exit_code)"
}