#!/usr/bin/expect
############################################################################
# Purpose: Test federation functionality
#          Test that submissions to clusters in federation get unique job id.
#
# 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.
#
# Output:  "TEST: #.#" followed by "SUCCESS" if test was successful, OR
#          "FAILURE: ..." otherwise with an explanation of the failure, OR
#          anything else indicates a failure mode that must be investigated.
############################################################################
# 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 test_id     "37.2"
set exit_code   0
set fed_name    "feda"

set eol "\r\n"

set timeout	5
print_header $test_id

#
# Check accounting config and bail if not found.
#
if { [test_account_storage] == 0 } {
	send_user "\nWARNING: This test can't be run without a usable AccountStorageType\n"
	exit 0
}

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

proc cleanup { } {
	global fed_name
	return [delete_federations $fed_name];
}

proc end_it { exit_code } {
	global test_id
	cleanup
	if {$exit_code == 0} {
		print_success $test_id
	}
	exit $exit_code
}

# submit jobs to each cluster and ensure that they return a unique job id
# job id should be cluster_id + job_id
# local_id:   bits 0-25
# cluster_id: bits26-31
proc test_fed_job_id { cname cid} {
	global fed_slurm_base number
	set job_id 0
	set rc 0

	#bits 0-25
	set max_local_id [expr 0x3FFFFFF]

	set my_sbatch "${fed_slurm_base}/$cname/bin/sbatch"
	set sbatch_pid [spawn $my_sbatch --wrap="hostname" -o/dev/null]
	expect {
		-re "Submitted batch job ($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}
		-re "error" {
			send_user "\nFAILURE: Job was not submitted\n"
			incr rc
		}
		timeout {
			send_user "\nFAILURE: sbatch is not responding\n"
			incr rc
		}
		eof {
			wait
		}
	}

	if {$job_id <= $max_local_id} {
		send_user "\nFAILURE: $cname didn't give back a federation jobid ($job_id)\n"
		incr rc
	}

	set local_id [expr $job_id & $max_local_id]
	set clust_id [expr $job_id >> 26]

	send_user "Fed JobID:$job_id Local JobID:$local_id Cluster ID:$clust_id\n"
	if {!$rc && ($clust_id != $cid)} {
		send_user "\nFAILURE: jobid($job_id) from $cname didn't give\
			correct partition id ($part_id != $cid)\n"
		incr rc
	}

	return $rc
}

# submit jobs to non-federated cluster and make sure the jobid is not a
# federated jobid.
proc test_loc_job_id { cname } {
	global fed_slurm_base number
	set job_id 0
	set rc 0

	#bits 0-25
	set max_local_id [expr 0x3FFFFFF]

	set my_sbatch "${fed_slurm_base}/$cname/bin/sbatch"
	set sbatch_pid [spawn $my_sbatch --wrap="hostname" -o/dev/null]
	expect {
		-re "Submitted batch job ($number)" {
			set job_id $expect_out(1,string)
			exp_continue
		}
		-re "error" {
			send_user "\nFAILURE: Job was not submitted\n"
			incr rc
		}
		timeout {
			send_user "\nFAILURE: sbatch is not responding\n"
			incr rc
		}
		eof {
			wait
		}
	}

	if {$job_id > $max_local_id} {
		send_user "\nFAILURE: $cname didn't give back a non-federation job id\n"
		incr rc
	}
	send_user "JobID:$job_id\n"

	return $rc
}

if {[test_federation_setup]} {
	send_user "\nWARNING: This test can't be run without fed_slurm_base,\
		fedc1, fedc2, fedc3 setup in globals.local.\n"
	exit 0
}

if {[test_all_up]} {
	exit 0
}

# Remove existing setup
if {[cleanup] != 0} {
	send_user "\nFAILURE: failed to cleanup\n"
	end_it 1
}

# add clusters to federation
if {[setup_federation $fed_name]} {
	send_user "\nFAILURE: failed to setup federation\n"
	end_it 1
}

# Get cluster/fed info from db
array set clusters [get_clusterfed_info $fed_name]

send_user "Verify federation ids are returned when cluster is part of federation\n"
if {[test_fed_job_id $fedc1 [dict get $clusters($fedc1) id]] ||
    [test_fed_job_id $fedc2 [dict get $clusters($fedc2) id]] ||
    [test_fed_job_id $fedc3 [dict get $clusters($fedc3) id]]} {
	end_it 1
}


# Verify that after clusters are removed from federation that they give normal
# job ids again.
if {[delete_federations $fed_name]} {
	send_user "\nFAILURE: failed delete federation $fed_name\n"
	end_it 1
}

sleep 2

send_user "Verify local ids after being removed from federation\n"
if {[test_loc_job_id $fedc1] ||
    [test_loc_job_id $fedc2] ||
    [test_loc_job_id $fedc3]} {
	end_it 1
}

end_it 0
