#!/usr/bin/env expect
############################################################################
# Purpose: Test federated submissions
#
# Reqs:    1. Using slurmdbd accounting storage type and is up
#          2. fed_slurm_base is defined in globals.local - set to directory that
#          has access to each federation configure (fedc1, fedc2, fedc3).
#          Eg.
#          fedr/slurm/ (src)
#          fedr/fed1/bin
#          fedr/fed1/sbin
#          fedr/fed1/etc
#          fedr/fed1/...
#          fedr/fed2/...
#          fedr/fed3/...
#          3. controllers are up and running.
############################################################################
# Copyright (C) 2016 SchedMD LLC.
# Written by Brian Christiansen <brian@schedmd.com>
#
# 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
source ./globals_accounting
source ./globals_federation

set fed_name    "feda"
set file_in     "test$test_id.in"
set user_name   ""
set srun_job_cnt 0

array set job_counts {}

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

if {[string compare [get_admin_level] "Administrator"]} {
	skip "This test can't be run without being an Accounting administrator. Use: sacctmgr mod user \$USER set admin=admin"
}

proc cancel_federation_jobs { } {
	global scancel user_name fedc1 fedc2 fedc3

	spawn $scancel -M$fedc1,$fedc2,$fedc3 --user $user_name
	expect {
		eof {
			wait
		}
	}
	sleep 5
}

proc cleanup { } {
	global scancel fed_name user_name bin_rm file_in fedc1 fedc2 fedc3
	global test_id bin_bash

	cancel_federation_jobs
	exec $bin_rm -f $file_in
	exec $bin_bash -c "$bin_rm -f test$test_id*.out"

	return [delete_federations $fed_name]
}

proc submit_fed_batch_job { cname expected_origin expected_sib extra_args expected_state } {
	global fed_slurm_base file_in node_count number squeue

	set submit_cluster ""
	set job_id 0
	set my_sbatch "${fed_slurm_base}/$cname/bin/sbatch"
	set command "$my_sbatch -N$node_count --exclusive --output=/dev/null --error=/dev/null -t3"
	if {$extra_args ne ""} {
		append command " $extra_args"
	}
	append command " $file_in"
	set sbatch_pid [spawn {*}$command]
	expect {
		-re "Submitted batch job ($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}
		-re "on cluster (\\S+)" {
			set submit_cluster $expect_out(1,string)
			exp_continue
		}
		timeout {
			slow_kill $sbatch_pid
			fail "sbatch not responding"
		}
		eof {
			wait
		}
	}
	if {$job_id == 0} {
		fail "Batch submit failure"
	}

	if {[string compare $expected_state ""]} {
		set run_cluster [wait_for_fed_job $job_id $expected_state ""]
		if {[string compare $run_cluster ""] == 0} {
			fail "Didn't find cluster with $expected_state job"
		}
	}

	return [validate_fed_job $cname $job_id $expected_origin $expected_sib $submit_cluster]
}

proc submit_fed_srun_job { cname expected_origin expected_sib extra_args check_env expected_state } {
	global fed_slurm_base file_in node_count number squeue test_id srun_job_cnt

	set submit_cluster ""
	set job_id 0
	set file_out "test${test_id}_${srun_job_cnt}.out"
	incr srun_job_cnt
	set my_srun "${fed_slurm_base}/$cname/bin/srun"
	set command "$my_srun -N$node_count --exclusive -t3 -o $file_out"
	if {$extra_args ne ""} {
		append command " $extra_args"
	}
	append command " $file_in"
	set srun_pid [spawn {*}$command]
	expect {
		-re "job ($number) queued and waiting for resources" {
			set job_id $expect_out(1,string)
		}
		timeout {
			fail "srun not responding"
		}
		eof {
			wait
		}
	}

	if {[string compare $expected_state ""]} {
		set run_cluster [wait_for_fed_job $job_id $expected_state ""]
		if {[string compare $run_cluster ""] == 0} {
			fail "Didn't find cluster with $expected_state job"
		}
	}

	if {$check_env} {
		if {[wait_for_file $file_out]} {
			fail "Failed waiting for file '$file_out'"
		}

		spawn grep SLURM_CLUSTER_NAME $file_out
		expect {
			-re "SLURM_CLUSTER_NAME=(\\S+)" {
				set submit_cluster $expect_out(1,string)
				if {$submit_cluster ne $expected_sib} {
					fail "SLURM_CLUSTER_NAME=$submit_cluster != $expected_sib"
				}
			}
			eof {
				wait
			}
		}
	}

	if {$job_id == 0} {
		fail "srun submit failure"
	}

	return [validate_fed_job $cname $job_id $expected_origin $expected_sib $submit_cluster]
}

proc submit_fed_salloc_job { cname expected_origin expected_sib extra_args check_env expected_state } {
	global fed_slurm_base file_in node_count number squeue test_id srun_job_cnt

	set submit_cluster ""
	set job_id 0
	set file_out "test${test_id}_${srun_job_cnt}.out"
	incr srun_job_cnt
	set my_salloc "${fed_slurm_base}/$cname/bin/salloc"
	set my_srun   "${fed_slurm_base}/$cname/bin/srun"
	set command "$my_salloc -N$node_count --exclusive -t3"
	if {$extra_args ne ""} {
		append command " $extra_args"
	}
	append command " $my_srun -o $file_out $file_in &"
	set salloc_pid [spawn {*}$command]
	expect {
		-re "salloc: Granted job allocation ($number)" {
			set job_id $expect_out(1,string)
		}
		-re "salloc: Pending job allocation ($number)" {
			set job_id $expect_out(1,string)
		}
		timeout {
			fail "salloc not responding"
		}
		eof {
			wait
		}
	}

	if {[string compare $expected_state ""]} {
		set run_cluster [wait_for_fed_job $job_id $expected_state ""]
		if {[string compare $run_cluster ""] == 0} {
			fail "Didn't find cluster with $expected_state job"
		}
	}

	if {$check_env} {
		if {[wait_for_file $file_out]} {
			fail "failed waiting for file '$file_out'"
		}

		spawn grep SLURM_CLUSTER_NAME $file_out
		expect {
			-re "SLURM_CLUSTER_NAME=(\\S+)" {
				set submit_cluster $expect_out(1,string)
				if {$submit_cluster ne $expected_sib} {
					fail "SLURM_CLUSTER_NAME=$submit_cluster != $expected_sib"
				}
			}
			eof {
				wait
			}
		}
	}

	if {$job_id == 0} {
		fail "salloc submit failure"
	}

	return [validate_fed_job $cname $job_id $expected_origin $expected_sib $submit_cluster]
}

proc validate_fed_job { cname job_id expected_origin expected_sib submit_cluster } {
	global fed_slurm_base file_in node_count number squeue

	set origin  ""
	set sibling ""

	sleep 3

	set my_squeue "${fed_slurm_base}/$cname/bin/squeue"
	if {$submit_cluster ne ""} {
		set my_squeue "${fed_slurm_base}/$submit_cluster/bin/squeue"
	}
	spawn $my_squeue -a --local --jobs=$job_id --noheader -Oorigin,siblingsactive
	expect {
		-re "(\\S+)\\s+(\\S+)" {
			set origin  $expect_out(1,string)
			set sibling $expect_out(2,string)
		}
	}

	log_info "origin:$origin sibling:$sibling"

	if {($expected_origin ne "") && ($origin ne $expected_origin)} {
		fail "origin:$origin != expected_origin:$expected_origin"
	}

	if {($expected_sib ne "") && ($sibling ne $expected_sib)} {
		fail "sibling:$sibling != expected_sib:$expected_sib"
	}


	# Verify that siblings have the job as well.
	foreach tmp_sib [split $sibling ","] {
		if {$tmp_sib eq  $origin} {
			continue
		}
		set my_squeue "${fed_slurm_base}/$tmp_sib/bin/squeue"
		spawn $my_squeue -a --local --jobs=$job_id --noheader -Oorigin,siblingsactive
		set match 0
		expect {
			-re "(\\S+)\\s+(\\S+)" {
				set match 1
				if {$origin ne $expect_out(1,string)} {
					fail "Origin not the same on $sibling"
				}
				if {$expect_out(2,string) ne "NA" && $sibling ne $expect_out(2,string)} {
					fail "Sibling not the same on $sibling"
				}
			}
			timeout {
				fail "$my_squeue not responding"
			}
			eof {
				wait
			}
		}

		if {!$match} {
			fail "Didn't find origin or sibling from job"
		}
	}

	return $sibling
}

proc validate_cluster_features { cname job_id expected_features } {
	global fed_slurm_base bin_bash bin_grep

	set my_scontrol "${fed_slurm_base}/$cname/bin/scontrol"
	set matches 0
	set scontrol_pid [spawn $bin_bash -c "$my_scontrol show jobid=$job_id | $bin_grep ClusterFeatures="]
	expect {
		"ClusterFeatures=$expected_features" {
			incr matches
		}
		timeout {
			slow_kill $scontrol_pid
			fail "scontrol not responding"
		}
		eof {
			wait
		}
	}
	if {$matches != 1} {
		fail "$matches didn't find ClusterFeatures=$expected_features on job"
	}
}

if {![check_federation_setup]} {
	skip "This test can't be run without fed_slurm_base, fedc1, fedc2, fedc3 setup in globals.local"
}

if {![check_federation_up]} {
	skip "This test can't be run without all clusters up"
}

set user_name [get_my_user_name]

# Remove existing setup
if {[cleanup] != 0} {
	fail "Failed to cleanup"
}

# add clusters to federation
if {[setup_federation $fed_name]} {
	fail "Failed to setup federation"
}

# get number of nodes per cluster
# devide by 2 to get 2 jobs per clusters
set node_count [expr [llength [get_nodes_by_state idle,alloc,comp]] / 2]

make_bash_script $file_in "env; $bin_sleep 300"


# Test needs to change to count that it gets two job per cluster instead of making sure that it packs on each. Because it's not guaranteed that the cluster will pack because the job gets submitted to each cluster and it's a free for all.
log_info "################################################################"
log_info "Test packing across clusters with sbatch"
log_info "################################################################"

unset job_counts

for {set i 0} {$i < 6} {incr i} {
	set ret_sib [submit_fed_batch_job $fedc1 $fedc1 "" "" RUNNING]
	incr job_counts($ret_sib)
}
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 "" "" PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc3) != 2} {
	fail "$fedc3 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "$fedc1,$fedc2,$fedc3 don't have a pending job on it"
}


log_info "################################################################"
log_info "Test packing across clusters with srun"
log_info "################################################################"
cancel_federation_jobs
unset job_counts

for {set i 0} {$i < 6} {incr i} {
	set ret_sib [submit_fed_srun_job $fedc1 $fedc1 "" "" 0 RUNNING]
	incr job_counts($ret_sib)
}
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 "" "" 0 PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc3) != 2} {
	fail "$fedc3 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "$fedc1,$fedc2,$fedc3 don't have a pending job on it"
}

log_info "################################################################"
log_info "Test packing across clusters with salloc"
log_info "################################################################"
cancel_federation_jobs
unset job_counts

for {set i 0} {$i < 6} {incr i} {
	set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 "" "" 0 RUNNING]
	incr job_counts($ret_sib)
}
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 "" "" 0 PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc3) != 2} {
	fail "$fedc3 doesn't have 2 jobs running on it"
}
if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "$fedc1,$fedc2,$fedc3 don't have a pending job on it"
}

log_info "################################################################"
log_info "Test -M<clusters> with federated jobs with sbatch"
log_info "################################################################"
cancel_federation_jobs
unset job_counts

# Submit job to only fed1
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc1 "-M$fedc1" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc1 "-M$fedc1" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc1 "-M$fedc1" PENDING]
incr job_counts($ret_sib)

# Submit job to only fed1,fed2
# Will go to fed2 since fed1 is full and third job should go to both
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc2 "-M$fedc1,$fedc2" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc2 "-M$fedc1,$fedc2" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 "$fedc1,$fedc2" "-M$fedc1,$fedc2" PENDING]
incr job_counts($ret_sib)

# Submit job to fed2,fed3.
# Should choose fed2 to be origin and submit
set ret_sib [submit_fed_batch_job $fedc1 $fedc2 $fedc3 "-M$fedc2,$fedc3" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc2 $fedc3 "-M$fedc2,$fedc3" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc2 "$fedc2,$fedc3" "-M$fedc2,$fedc3" PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 3} {
	fail "$fedc1 doesn't have 3 jobs running on it"
}

if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc1,$fedc2) != 1} {
	fail "$fedc1,$fedc2 doesn't have 1 jobs running on it"
}

if {$job_counts($fedc3) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc2,$fedc3) != 1} {
	fail "$fedc2,$fedc3 doesn't have 1 jobs running on it"
}

log_info "################################################################"
log_info "Test -M<clusters> with federated jobs with srun"
log_info "################################################################"
cancel_federation_jobs
unset job_counts

# Submit job to only fed1
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc1 "-M$fedc1" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc1 "-M$fedc1" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc1 "-M$fedc1" 0 PENDING]
incr job_counts($ret_sib)

# Submit job to only fed1,fed2
# Will go to fed2 since fed1 is full and third job should go to both
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc2 "-M$fedc1,$fedc2" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc2 "-M$fedc1,$fedc2" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 "$fedc1,$fedc2" "-M$fedc1,$fedc2" 0 PENDING]
incr job_counts($ret_sib)

# Submit job to fed2,fed3.
# Should choose fed2 to be origin and submit
set ret_sib [submit_fed_srun_job $fedc1 $fedc2 $fedc3 "-M$fedc2,$fedc3" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc2 $fedc3 "-M$fedc2,$fedc3" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc2 $fedc2 "$fedc2,$fedc3" "-M$fedc2,$fedc3" 0 PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 3} {
	fail "$fedc1 doesn't have 3 jobs running on it"
}

if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc1,$fedc2) != 1} {
	fail "$fedc1,$fedc2 doesn't have 1 jobs running on it"
}

if {$job_counts($fedc3) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc2,$fedc3) != 1} {
	fail "$fedc2,$fedc3 doesn't have 1 jobs running on it"
}


log_info "################################################################"
log_info "Test -M<clusters> with federated jobs with salloc"
log_info "################################################################"
cancel_federation_jobs
unset job_counts

# Submit job to only fed1
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc1 "-M$fedc1" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc1 "-M$fedc1" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc1 "-M$fedc1" 0 PENDING]
incr job_counts($ret_sib)

# Submit job to only fed1,fed2
# Will go to fed2 since fed1 is full and third job should go to both
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc2 "-M$fedc1,$fedc2" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc2 "-M$fedc1,$fedc2" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 "$fedc1,$fedc2" "-M$fedc1,$fedc2" 0 PENDING]
incr job_counts($ret_sib)

# Submit job to fed2,fed3.
# Should choose fed2 to be origin and submit
set ret_sib [submit_fed_salloc_job $fedc1 $fedc2 $fedc3 "-M$fedc2,$fedc3" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc2 $fedc3 "-M$fedc2,$fedc3" 1 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc2 $fedc2 "$fedc2,$fedc3" "-M$fedc2,$fedc3" 0 PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 3} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc1,$fedc2) != 1} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc3) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc2,$fedc3) != 1} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}


log_info "################################################################"
log_info "Test jobs with cluster features"
log_info "################################################################"
cancel_federation_jobs

set matches 0
set my_pid [spawn $sacctmgr -i modify cluster $fedc1 set features=fa]
expect {
	-re "Setting$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+Feature\\s+=\\s+fa" {
		incr matches
		exp_continue
	}
	-re "Modified cluster...$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+$fedc1$eol" {
		incr matches
		exp_continue
	}
	timeout {
		slow_kill $my_pid
		fail "sacctmgr mod not responding"
	}
	eof {
		wait
	}
}
if {$matches != 4} {
	fail "Unexpected error. Got $matches"
}

set matches 0
set my_pid [spawn $sacctmgr -i modify cluster $fedc2 set features=fb]
expect {
	-re "Setting$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+Feature\\s+=\\s+fb" {
		incr matches
		exp_continue
	}
	-re "Modified cluster...$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+$fedc2$eol" {
		incr matches
		exp_continue
	}
	timeout {
		slow_kill $my_pid
		fail "sacctmgr mod not responding"
	}
	eof {
		wait
	}
}
if {$matches != 4} {
	fail "Unexpected error. Got $matches"
}

set matches 0
set my_pid [spawn $sacctmgr -i modify cluster $fedc3 set features=fc]
expect {
	-re "Setting$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+Feature\\s+=\\s+fc" {
		incr matches
		exp_continue
	}
	-re "Modified cluster...$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+$fedc3$eol" {
		incr matches
		exp_continue
	}
	timeout {
		slow_kill $my_pid
		fail "sacctmgr mod not responding"
	}
	eof {
		wait
	}
}
if {$matches != 4} {
	fail "Unexpected error. Got $matches"
}

log_info "################################################################"
log_info "Test sbatch jobs with cluster features"
log_info "################################################################"
unset job_counts

set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc3 "--cluster-constraint=fc" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc3 "--cluster-constraint=fc" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc2 "--cluster-constraint=fb" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc2 "--cluster-constraint=fb" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc1 "--cluster-constraint=fa" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 $fedc1 "--cluster-constraint=fa" RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_batch_job $fedc1 $fedc1 "$fedc1,$fedc2,$fedc3" "--cluster-constraint=fa,fb,fc" PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc3) != 2} {
	fail "$fedc3 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "$fedc1,$fedc2,$fedc3 doesn't have 2 jobs running on it"
}

log_info "################################################################"
log_info "Test salloc jobs with cluster features"
log_info "################################################################"
cancel_federation_jobs
unset job_counts

set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc3 "--cluster-constraint=fc" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc3 "--cluster-constraint=fc" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc2 "--cluster-constraint=fb" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc2 "--cluster-constraint=fb" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc1 "--cluster-constraint=fa" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 $fedc1 "--cluster-constraint=fa" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_salloc_job $fedc1 $fedc1 "$fedc1,$fedc2,$fedc3" "--cluster-constraint=fa,fb,fc" 0 PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc3) != 2} {
	fail "$fedc3 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "$fedc1,$fedc2,$fedc3 doesn't have 2 jobs running on it"
}

log_info "################################################################"
log_info "Test srun jobs with cluster features"
log_info "################################################################"
cancel_federation_jobs
unset job_counts

set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc3 "--cluster-constraint=fc" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc3 "--cluster-constraint=fc" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc2 "--cluster-constraint=fb" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc2 "--cluster-constraint=fb" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc1 "--cluster-constraint=fa" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 $fedc1 "--cluster-constraint=fa" 0 RUNNING]
incr job_counts($ret_sib)
set ret_sib [submit_fed_srun_job $fedc1 $fedc1 "$fedc1,$fedc2,$fedc3" "--cluster-constraint=fa,fb,fc" 0 PENDING]
incr job_counts($ret_sib)

parray job_counts

if {$job_counts($fedc1) != 2} {
	fail "$fedc1 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc2) != 2} {
	fail "$fedc2 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc3) != 2} {
	fail "$fedc3 doesn't have 2 jobs running on it"
}

if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "$fedc1,$fedc2,$fedc3 doesn't have 2 jobs running on it"
}

log_info "################################################################"
log_info "Test jobs with invalid cluster features"
log_info "################################################################"
set matches 0
set my_sbatch "${fed_slurm_base}/$fedc1/bin/sbatch"
set sbatch_pid [spawn $my_sbatch --output=/dev/null --error=/dev/null -n1 -t3 --cluster-constraint=invalid $file_in]
expect {
	"sbatch: error: Batch job submission failed: Invalid cluster feature specification" {
		incr matches
	}
	timeout {
		slow_kill $sbatch_pid
		fail "sbatch not responding"
	}
	eof {
		wait
	}
}
if {$matches != 1} {
	fail "Invalid cluster features sbatch didn't fail. Got $matches"
}
