#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Verify node configuration specification (--constraint option).
############################################################################
# Copyright (C) SchedMD LLC.
# Copyright (C) 2002-2006 The Regents of the University of California.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Morris Jette
#
# 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 test_part   "$test_name\_part"
set feat0       "$test_name\_feat0"
set feat1       "$test_name\_feat1"
set feat2       "$test_name\_feat2"
set feat3       "$test_name\_feat3"
set feat4       "$test_name\_feat4"
set wrap_cmd    "--wrap='$bin_sleep 5'"
array set nodes {}

proc cleanup {} {
	global scontrol test_part

	# Remove test partition
	wait_for_part_done $test_part
	run_command "$scontrol delete partition=$test_part"

	# Reset node features
	reconfigure
}

#
# Prerequisites
#
if {![is_super_user] || [get_config_param "NodeFeaturesPlugins"] ne "(null)"} {
	skip "Test needs super user and not NodeFeaturesPlugins"
}

set nodelist [get_nodes_by_request "-N4 -t1"]
if {[llength $nodelist] != 4} {
	skip "Tests needs to be able to submit up to 4 nodes"
}
set nodes(0) [lindex $nodelist 0]
set nodes(1) [lindex $nodelist 1]
set nodes(2) [lindex $nodelist 2]
set nodes(3) [lindex $nodelist 3]

proc test_constrain_invalid {nnode constrain} {
	global sbatch scontrol test_part wrap_cmd

	# Submit the job
	set test_job 0
	set result [run_command -none "$sbatch -C \"$constrain\" -p$test_part -N$nnode -t1 -o /dev/null $wrap_cmd"]
	set output [dict get $result output]
	set rc     [dict get $result exit_code]
	regexp {Submitted batch job (\d+)} $output - test_job

	# Check expected output
	subtest {$test_job == 0} "Job should have been rejected"
	subtest {[regexp "Invalid feature specification" $output]} "Job should be rejected due invalid feature"
	cancel_job $test_job
}

proc test_constrain {nnode constrain mandatory_nodes {posible_nodes ""} {excluded_nodes ""}} {
	global sbatch scontrol test_part wrap_cmd

	# Submit the job
	set test_job 0
	set result [run_command -none "$sbatch -C \"$constrain\" -p$test_part -N$nnode -t1 -o /dev/null $wrap_cmd"]
	set output [dict get $result output]
	set rc     [dict get $result exit_code]
	regexp {Submitted batch job (\d+)} $output - test_job

	# Check expected output
	if {![llength $mandatory_nodes] && ![llength $posible_nodes]} {
		subtest {$test_job == 0} "Job should have been rejected"
		subtest {[regexp "Requested node configuration is not available" $output]} "Job should be rejected due configuration not available"
	} else {
		if {[subtest {$test_job != 0} "Job should be submitted"]} {
			wait_for_job -fail $test_job "RUNNING"

			# Check that job that the job used the correct nodes
			set hostslist [get_job_param $test_job "NodeList"]
			set nodeslist [run_command_output -fail "$scontrol show hostnames $hostslist"]
			# TODO: make nodeslist an actual list
			subtest {[llength $nodeslist] == $nnode} "Job's NodeList should contain $nnode nodes" "[llength $nodeslist]"
			foreach node $nodeslist {
				subtest {[lsearch $posible_nodes $node] != -1 || [lsearch $mandatory_nodes $node] != -1} "Job's node ($node) should be in the expected list" "$posible_nodes $mandatory_nodes"
				subtest {[lsearch $excluded_nodes $node] == -1 } "Job's node ($node) should NOT be in the excluded list" "$excluded_nodes"
			}
			foreach node $mandatory_nodes {
				subtest {[lsearch $nodeslist $node] != -1} "Node ($node) should be in the jobs NodeList" "$nodeslist"
			}
		}
	}
	cancel_job $test_job
}

proc test_constrain_batch {nnode constrain constrain_batch expected_batch_node} {
	global sbatch scontrol test_part wrap_cmd

	# Submit the job
	set test_job 0
	set result [run_command -none "$sbatch -C \"$constrain\" --batch=$constrain_batch -p$test_part -N$nnode -t1 -o /dev/null $wrap_cmd"]
	set output [dict get $result output]
	set rc     [dict get $result exit_code]
	regexp {Submitted batch job (\d+)} $output - test_job

	if {[subtest {$test_job != 0} "Job should be submitted"]} {
		wait_for_job -fail $test_job "RUNNING"
		# Check that job that the job used the correct nodes
		set batch_host [get_job_param $test_job "BatchHost"]
		subtest {$batch_host == $expected_batch_node} "Verify job's expected batch node" "$batch_host != $expected_batch_node"
	}
	cancel_job $test_job
}

proc test_constrain_9546 {} {
	global test_name feat0 feat1 bin_sleep test_part

	set job_id_1 [submit_job -fail  "--exclusive  -J $test_name -C $feat0 -p $test_part -o /dev/null --wrap '$bin_sleep 1000'" ]
	set job_id_2 [submit_job -fail  "--exclusive  -J $test_name -C $feat0 -p $test_part -o /dev/null --wrap '$bin_sleep 1000'" ]

	subtest {![wait_for_job $job_id_1 RUNNING]} "Job1 ($job_id_1) with \"-C $feat0\" should be started"
	subtest {![wait_for_job $job_id_2 RUNNING]} "Job2 ($job_id_2) with \"-C $feat0\" should be started"

	set job_id_3 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat1|$feat0\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]
	set job_id_4 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat1|$feat0\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]

	subtest {![wait_for_job $job_id_3 DONE]} "Job3 ($job_id_3) with \"-C '\[$feat1|$feat0\]'\" should be completed"
	subtest {![wait_for_job $job_id_4 DONE]} "Job4 ($job_id_4) with \"-C '\[$feat1|$feat0\]'\" should be completed"

	set job_id_3 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat0|$feat1\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]
	set job_id_4 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat0|$feat1\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]

	subtest {![wait_for_job $job_id_3 DONE]} "Job3 ($job_id_3) with \"-C '\[$feat0|$feat1\]'\" should be completed"
	subtest {![wait_for_job $job_id_4 DONE]} "Job4 ($job_id_4) with \"-C '\[$feat0|$feat1\]'\" should be completed"

	cancel_job [list $job_id_1 $job_id_2]

	set job_id_1 [submit_job -fail  "--exclusive  -J $test_name -C $feat1 -p $test_part -o /dev/null --wrap '$bin_sleep 1000'" ]
	set job_id_2 [submit_job -fail  "--exclusive  -J $test_name -C $feat1 -p $test_part -o /dev/null --wrap '$bin_sleep 1000'" ]

	subtest {![wait_for_job $job_id_1 RUNNING]} "Job1 ($job_id_1) with \"-C $feat1\" should be started"
	subtest {![wait_for_job $job_id_2 RUNNING]} "Job2 ($job_id_2) with \"-C $feat1\" should be started"

	set job_id_3 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat1|$feat0\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]
	set job_id_4 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat1|$feat0\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]

	subtest {![wait_for_job $job_id_3 DONE]} "Job3 ($job_id_3) with \"-C '\[$feat1|$feat0\]'\" should be completed"
	subtest {![wait_for_job $job_id_4 DONE]} "Job4 ($job_id_4) with \"-C '\[$feat1|$feat0\]'\" should be completed"

	set job_id_3 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat0|$feat1\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]
	set job_id_4 [submit_job -fail  "--exclusive  -J $test_name -C \'\[$feat0|$feat1\]\' -p $test_part -o /dev/null --wrap 'exit 0'" ]

	subtest {![wait_for_job $job_id_3 DONE]} "Job3 ($job_id_3) with \"-C '\[$feat0|$feat1\]'\" should be completed"
	subtest {![wait_for_job $job_id_4 DONE]} "Job4 ($job_id_4) with \"-C '\[$feat0|$feat1\]'\" should be completed"
}

proc set_node_feature {node_name new_feature} {
	global scontrol
	run_command -fail "$scontrol update node=$node_name  AvailableFeatures=$new_feature ActiveFeatures=$new_feature"
}


proc setup_4_features {} {
	global scontrol test_part
	global nodes nodelist feat0 feat1 feat2 feat3 feat4

	log_info "Setup partition $test_part with 4 nodes with 4 features"
	set_node_feature $nodes(0) $feat0
	set_node_feature $nodes(1) $feat1
	set_node_feature $nodes(2) $feat0,$feat2,$feat4
	set_node_feature $nodes(3) $feat0,$feat2,$feat3,$feat4

	run_command -fail "$scontrol create partition=$test_part nodes=[join $nodelist ,]"
}
proc setup_2_features {} {
	global scontrol test_part
	global nodes nodelist feat0 feat1 feat2 feat3 feat4

	log_info "Setup partition $test_part with 4 nodes with 2 features"
	set_node_feature $nodes(0) $feat0
	set_node_feature $nodes(1) $feat0
	set_node_feature $nodes(2) $feat1
	set_node_feature $nodes(3) $feat1

	run_command -fail "$scontrol create partition=$test_part nodes=[join $nodelist ,]"
}

#
# Main setup of 4 features on 4 nodes
#
setup_4_features

#
# Main/Generic testprocs
#

# Test feature count logic with 2 nodes
testproc test_constrain 2 $feat2*2 [list $nodes(2) $nodes(3)]

# Test feature count logic with 3 nodes
testproc test_constrain 3 $feat2*2 [list $nodes(2) $nodes(3)] [list $nodes(0) $nodes(1)]

# Test feature count logic with 3 nodes (This is expected to fail)
testproc test_constrain 3 $feat2*3 [list]

# Test AND with 1 node
testproc test_constrain 1 $feat2&$feat3 [list $nodes(3)]

# Test AND with 1 node (This is expected to fail)
testproc test_constrain 1 $feat1&$feat3 [list]

# Test eXclusive AND (XAND) with 3 nodes
testproc test_constrain 3 "\[($feat0&$feat2)*2&$feat1*1\]" [list $nodes(1) $nodes(2) $nodes(3)]

# Test eXclusive AND (XAND) with 3 nodes (This is expected to fail)
testproc test_constrain 3 "\[($feat0&$feat3)*2&$feat1*2\]" [list]

# Test eXclusive OR (XOR) with 2 nodes
testproc test_constrain 2 "\[$feat0|$feat1\]" "" [list $nodes(0) $nodes(2) $nodes(3)] [list $nodes(1)]
# testproc sub_job_xor 2 0

# Test eXclusive OR (XOR) with 3 nodes
testproc test_constrain 3 "\[$feat0|$feat1\]" "" [list $nodes(0) $nodes(2) $nodes(3)] [list $nodes(1)]

# Test eXclusive OR (XOR) with 4 nodes (This is expected to fail)
testproc test_constrain 4 "\[$feat0|$feat1\]" [list]

# Test OR with 4 nodes
testproc test_constrain 4 "$feat0|$feat1" [list $nodes(0) $nodes(1) $nodes(2) $nodes(3)]

# Test OR within parenthesis
testproc test_constrain 4 "($feat0|$feat1)" "" [list $nodes(0) $nodes(1) $nodes(2) $nodes(3)]
testproc test_constrain 3 "($feat0|$feat2)" [list $nodes(0) $nodes(2) $nodes(3)] [list $nodes(1)]
testproc test_constrain 2 "($feat1|$feat3)" [list $nodes(1) $nodes(3)] [list $nodes(0) $nodes(2)]
testproc test_constrain 2 "$feat0&($feat2|$feat3)" [list $nodes(2) $nodes(3)] [list $nodes(0) $nodes(1)]

# Test mixed AND and XOR with 2 nodes
testproc test_constrain 2 "$feat4&\[$feat1|$feat0\]" "" [list $nodes(0) $nodes(2) $nodes(3)] [list $nodes(1)]

# Test mixed AND and XOR with 3 nodes (This is expected to fail)
testproc test_constrain 3 "$feat4&\[$feat1|$feat0\]" [list]

#
# Special testprocs
#
# Specific test for batch host
testproc test_constrain_batch 3 "$feat0|$feat3" $feat3 [list $nodes(3)]
run_following_testprocs

# Specific test invalid constraint requirement
testproc test_constrain_invalid 1 "invalid,constraint"

# Specific test for bug 9546 about xor/xand (with specific setup)
cleanup
setup_2_features
testproc test_constrain_9546