#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Test srun stdout/err labelling combined with file template
#          options with %t and %N.
############################################################################
# Copyright (C) 2009-2010 Lawrence Livermore National Security.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Dave Bremer <dbremer@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_out_n       "test$test_id.n.%n.output"
set file_out_t       "test$test_id.t.%t.output"
set job_id           0
set node_count       0
set task_count       0
set task_id          0
set node_id          0
set file_cnt         0
set file_out_t_glob  ""
set file_out_n_glob  ""

if {[test_front_end]} {
    skip "This test is incompatible with front-end systems"
}

#
# Spawn a program that generates "task_id" (%t) in stdout file names
# and confirm they are created
#
for {set task_id 0} {$task_id < $node_count} {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 -l --output=$file_out_t -N 1-10 -v -t1 $bin_echo hello]
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 node_count 0
spawn $squeue -tall -j $job_id -o "%i %D"
expect {
	-re "$job_id ($number)" {
		set node_count $expect_out(1,string)
		exp_continue
	}
	eof {
		wait
	}
}
if {$node_count == 0} {
	cancel_job $job_id
	fail "Did not get node_count"
}

set file_cnt 0
for {set task_id 0} {$task_id < $node_count} {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
		break
	} else {
		set test_task_id -1
		spawn $bin_cat $file_out_t_glob
		expect {
			-re "($number): *hello" {
				set test_task_id $expect_out(1,string)
				exp_continue
			}
			eof {
				wait
			}
		}
		if {$task_id != $test_task_id} {
			log_error "File $file_out_t_glob was not labelled.  $task_id != $test_task_id"
			set exit_code 1
		} else {
			exec $bin_rm -f $file_out_t_glob
		}
		incr file_cnt
	}
}
if {$file_cnt != $node_count} {
	log_error "File format of %t in stdout failed"
	set exit_code 1
}
fail_on_error "Test failed due to previous errors"
for {set node_id 0} {$node_id < $node_count} {incr node_id} {
	set file_out_n_glob  "test$test_id.n.$node_id.output"
	exec $bin_rm -f $file_out_n_glob
}

if {[test_front_end] != 0} {
	skip "Additional tests are incompatible with front-end systems"
}

#
# Spawn a program that generates "node_id" (%n) in stdout file names
# and confirm they are created
#
set task_count [expr $node_count * 2]
set timeout $max_job_delay
set srun_pid [spawn $srun -l --output=$file_out_n -N $node_count -n $task_count -O -v -t1 $bin_echo hello]
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"
}

# We only test that some task ran on every node, we can't check the task ID
# since task distribution is dependent upon resource allocation which may
# not be homogeneous across the nodes.
for {set node_id 0} {$node_id < $node_count} {incr node_id} {
	set file_out_n_glob  "test$test_id.n.$node_id.output"
	if {[wait_for_file $file_out_n_glob] != 0} {
		if {$node_id == 1} {
			log_debug "This failure is expected on front-end configurations"
		}
		set exit_code 1
		break
	} else {
		set task_found false

		spawn $bin_cat $file_out_n_glob
		expect {
			-re "($number): *hello" {
				set task_found true
				exp_continue
			}
			eof {
				wait
			}
		}
		if {!$task_found} {
			log_error "File $file_out_n_glob was not labelled correctly"
			set exit_code 1
		} else {
			exec $bin_rm -f $file_out_n_glob
		}
		incr file_cnt
	}
}

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