#!/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) SchedMD LLC.
#
# 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_dir/input"
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 {[get_admin_level] ne "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 fedc1 fedc2 fedc3
	global test_name bin_bash

	cancel_federation_jobs
	exec $bin_bash -c "$bin_rm -f $test_name*.out"

	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"
	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 {
			fail "sbatch not responding"
		}
		eof {
			wait
		}
	}
	if {$job_id == 0} {
		fail "Batch submit failure"
	}

	if {$expected_state ne ""} {
		set run_cluster [wait_for_fed_job $job_id $expected_state ""]
		if {$run_cluster eq ""} {
			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 {$expected_state ne ""} {
		set run_cluster [wait_for_fed_job $job_id $expected_state ""]
		if {$run_cluster eq ""} {
			fail "Didn't find cluster with $expected_state job"
		}
	}

	if {$check_env} {
		wait_for_file -fail $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 {$expected_state ne ""} {
		set run_cluster [wait_for_fed_job $job_id $expected_state ""]
		if {$run_cluster eq ""} {
			fail "Didn't find cluster with $expected_state job"
		}
	}

	if {$check_env} {
		wait_for_file -fail $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
	spawn $bin_bash -c "$my_scontrol show jobid=$job_id | $bin_grep ClusterFeatures="
	expect {
		"ClusterFeatures=$expected_features" {
			incr matches
		}
		timeout {
			fail "scontrol not responding"
		}
		eof {
			wait
		}
	}
	if {$matches != 1} {
		fail "Didn't find ClusterFeatures=$expected_features on job ($matches != 1)"
	}
}

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
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 "[default_partition] --local"]] / 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 "Cluster 1 ($fedc1) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc2) != 2} {
	fail "Cluster 2 ($fedc2) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc3) != 2} {
	fail "Cluster 3 ($fedc3) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "Clusters ($fedc1,$fedc2,$fedc3) don't have a pending job on them"
}


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 "Cluster 1 ($fedc1) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc2) != 2} {
	fail "Cluster 2 ($fedc2) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc3) != 2} {
	fail "Cluster 3 ($fedc3) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "Clusters ($fedc1,$fedc2,$fedc3) don't have a pending job on them"
}

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 "Cluster 1 ($fedc1) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc2) != 2} {
	fail "Cluster 2 ($fedc2) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc3) != 2} {
	fail "Cluster 3 ($fedc3) doesn't have 2 jobs running on it"
}
if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "Clusters ($fedc1,$fedc2,$fedc3) don't have a pending job on them"
}

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 "Cluster 1 ($fedc1) doesn't have 3 jobs running on it"
}

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

if {$job_counts($fedc1,$fedc2) != 1} {
	fail "Clusters 1 and 2 ($fedc1,$fedc2) don't have 1 job running on them"
}

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

if {$job_counts($fedc2,$fedc3) != 1} {
	fail "Clusters 2 and 3 ($fedc2,$fedc3) don't have 1 jobs running on them"
}

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 "Cluster 1 ($fedc1) doesn't have 3 jobs running on it"
}

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

if {$job_counts($fedc1,$fedc2) != 1} {
	fail "Clusters 1 and 2 ($fedc1,$fedc2) don't have 1 job running on them"
}

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

if {$job_counts($fedc2,$fedc3) != 1} {
	fail "Clusters 2 and 3 ($fedc2,$fedc3) don't have 1 job running on them"
}


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 "Cluster 1 ($fedc1) doesn't have 2 jobs running on it"
}

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

if {$job_counts($fedc1,$fedc2) != 1} {
	fail "Clusters 1 and 2 ($fedc1,$fedc2) don't have 1 job running on them"
}

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

if {$job_counts($fedc2,$fedc3) != 1} {
	fail "Clusters 2 and 3 ($fedc2,$fedc3) don't have 1 job running on them"
}


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

set matches 0
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 {
		fail "sacctmgr mod not responding"
	}
	eof {
		wait
	}
}
if {$matches != 4} {
	fail "Unexpected error (got $matches)"
}

set matches 0
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 {
		fail "sacctmgr mod not responding"
	}
	eof {
		wait
	}
}
if {$matches != 4} {
	fail "Unexpected error (got $matches)"
}

set matches 0
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 {
		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 "Cluster 1 ($fedc1) doesn't have 2 jobs running on it"
}

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

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

if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "Clusters ($fedc1,$fedc2,$fedc3) don't have 1 job running on them"
}

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 "Cluster 1 ($fedc1) doesn't have 2 jobs running on it"
}

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

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

if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "Clusters ($fedc1,$fedc2,$fedc3) don't have 1 job running on them"
}

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 "Cluster 1 ($fedc1) doesn't have 2 jobs running on it"
}

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

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

if {$job_counts($fedc1,$fedc2,$fedc3) != 1} {
	fail "Clusters ($fedc1,$fedc2,$fedc3) don't have 1 job running on them"
}

log_info "################################################################"
log_info "Test jobs with invalid cluster features"
log_info "################################################################"
set matches 0
set my_sbatch "${fed_slurm_base}/$fedc1/bin/sbatch"
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 {
		fail "sbatch not responding"
	}
	eof {
		wait
	}
}
if {$matches != 1} {
	fail "Invalid cluster features sbatch didn't fail (got $matches)"
}
