#!/usr/bin/expect
############################################################################
# Purpose: Test federation functionality
#          Test scontrol show federations
#	   Test that persistent connections are made
#
# 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.
#          federation/slurm/ (src)
#          federation/fed1/bin
#          federation/fed1/sbin
#          federation/fed1/etc
#          federation/fed1/...
#          federation/fed2/...
#          federation/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.3"
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
}

proc test_fed_status {cname fed_flags cluster_list} {
	# Match db info with what cluster has.
	global fed_slurm_base fed_name eol
	array set clusters $cluster_list
	set matches 0
	set lines 0
	set my_scontrol "${fed_slurm_base}/$cname/bin/scontrol"
	set my_pid [spawn $my_scontrol show federation]
	set expected_matches [array size clusters]
	if {$expected_matches} {
		# Match the fedname line
		incr expected_matches
	}
	expect {
		-re "$eol" {
			incr lines
			set line $expect_out(buffer)
			send_user "line: $line\n"

			if {[regexp "Federation: $fed_name" $line match]} {
				send_user "matched: $match\n"
				incr matches
			} elseif {[regexp {Self:\s+(\S+):(\S+):(\d+) ID:(\d+) FedState:(\S*) Features:(\S*)} \
					  $line match name host port id state features]} {
				send_user "matched: $match\n"
				if {$expected_matches &&
				    ![string compare [dict get $clusters($name) host]     $host]  &&
				    ![string compare [dict get $clusters($name) state]    $state] &&
				    ![string compare [dict get $clusters($name) features] $features]  &&
				    ![string compare [dict get $clusters($name) conn]     Self]  &&
				    [dict get $clusters($name) port]   == $port &&
				    [dict get $clusters($name) id]     == $id} {
					send_user "matched self: $name\n"
					incr matches
				}
			} elseif {[regexp {Sibling:\s+(\S+):(\S+):(\d+) ID:(\d+) FedState:(\S*) Features:(\S*) PersistConnSend/Recv:(\S+)} \
					  $line match name host port id state features conn]} {
				send_user "matched: $match\n"
				if {$expected_matches &&
				    ![string compare [dict get $clusters($name) host]     $host]  &&
				    ![string compare [dict get $clusters($name) state]    $state] &&
				    ![string compare [dict get $clusters($name) features] $features] &&
				    ![string compare [dict get $clusters($name) conn]     $conn]  &&
				    [dict get $clusters($name) port]   == $port &&
				    [dict get $clusters($name) id]     == $id} {
					send_user "matched sibling: $name\n"
					incr matches
				}
			}
			exp_continue
		}
		timeout {
			send_user "\nFAILURE: scontrol not responding\n"
			slow_kill $my_pid
			end_it 1
		}
		eof {
			wait
		}
	}

	if {$lines != $expected_matches} {
		send_user "FAILURE: line count and expected matches don't match: $lines != $expected_matches.\n"
		end_it 1
	}

	if {$matches != $expected_matches} {
		send_user "FAILURE: failed to match $matches != $expected_matches.\n"
		end_it 1
	}
}

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]

set fed_flags "None"

send_user "\n\ntest from $fedc1\n"
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
dict set clusters($fedc3) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]

send_user "\n\ntest from $fedc2\n"
dict set clusters($fedc1) conn "Yes/Yes"
dict set clusters($fedc2) conn Self
dict set clusters($fedc3) conn "Yes/Yes"
test_fed_status $fedc2 $fed_flags [array get clusters]

send_user "\n\ntest from $fedc3\n"
dict set clusters($fedc1) conn "Yes/Yes"
dict set clusters($fedc2) conn "Yes/Yes"
dict set clusters($fedc3) conn Self
test_fed_status $fedc3 $fed_flags [array get clusters]

send_user "\n\nremove $fedc3 and test from $fedc1 and $fedc2\n"
if {[remove_cluster_from_fed $fedc3 $fed_name]} {
	end_it 1
}
array unset clusters $fedc3
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]

dict set clusters($fedc1) conn "Yes/Yes"
dict set clusters($fedc2) conn Self
test_fed_status $fedc2 $fed_flags [array get clusters]

send_user "\n\nremove $fedc1 and test from $fedc2\n"
if {[remove_cluster_from_fed $fedc1 $fed_name]} {
	end_it 1
}
array unset clusters $fedc1
dict set clusters($fedc2) conn Self
test_fed_status $fedc2 $fed_flags [array get clusters]

send_user "\n\nremove $fedc2 and test from $fedc2\n"
array unset clusters
if {[remove_cluster_from_fed $fedc2 $fed_name]} {
	end_it 1
}
test_fed_status $fedc2 $fed_flags [array get clusters]

send_user "\n\nadd $fedc1 and test from $fedc1\n"
if {[add_cluster_to_fed $fedc1 $fed_name]} {
	end_it 1
}
array set clusters [get_clusterfed_info $fed_name]
dict set clusters($fedc1) conn Self
test_fed_status $fedc1 $fed_flags [array get clusters]

send_user "\n\nadd $fedc2 and test from $fedc1\n"
if {[add_cluster_to_fed $fedc2 $fed_name]} {
	end_it 1
}
sleep 5
array set clusters [get_clusterfed_info $fed_name]
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]

send_user "\n\nadd $fedc3 and test from $fedc1\n"
if {[add_cluster_to_fed $fedc3 $fed_name]} {
	end_it 1
}
sleep 5
array set clusters [get_clusterfed_info $fed_name]
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
dict set clusters($fedc3) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]


#send_user "\n\nTest federation flags (LLC)\n"
#set matches 0
#set my_pid [spawn $sacctmgr -i modify federation $fed_name set flags+=LLC]
#expect {
#	-re "Setting$eol" {
#		incr matches
#		exp_continue
#	}
#	-re "^\\s+Flags\\s+\\+=\\s+LLC$eol" {
#		incr matches
#		exp_continue
#	}
#	-re "^\\s+Modified federation...$eol" {
#		incr matches
#		exp_continue
#	}
#	-re "^\\s+$fed_name$eol" {
#		incr matches
#		exp_continue
#	}
#	timeout {
#		send_user "\nFAILURE: sacctmgr add not responding\n"
#		slow_kill $my_pid
#		set exit_code 1
#	}
#	eof {
#		wait
#	}
#}
#if {$exit_code || $matches != 4} {
#	send_user "$matches FAILURE: setting LLC flag.\n"
#	end_it 1
#}
#
#set fed_flags "LLC"
#array set clusters [get_clusterfed_info $fed_name]
#dict set clusters($fedc1) conn Self
#dict set clusters($fedc2) conn "Yes/Yes"
#dict set clusters($fedc3) conn "Yes/Yes"
#test_fed_status $fedc1 $fed_flags [array get clusters]
#
#
#send_user "\n\nTest removal of federation flags (LLC)\n"
#set matches 0
#set my_pid [spawn $sacctmgr -i modify federation $fed_name set flags-=LLC]
#expect {
#	-re "Setting$eol" {
#		incr matches
#		exp_continue
#	}
#	-re "^\\s+Flags\\s+\\-=\\s+LLC$eol" {
#		incr matches
#		exp_continue
#	}
#	-re "^\\s+Modified federation...$eol" {
#		incr matches
#		exp_continue
#	}
#	-re "^\\s+$fed_name$eol" {
#		incr matches
#		exp_continue
#	}
#	timeout {
#		send_user "\nFAILURE: sacctmgr add not responding\n"
#		slow_kill $my_pid
#		set exit_code 1
#	}
#	eof {
#		wait
#	}
#}
#if {$exit_code || $matches != 4} {
#	send_user "$matches FAILURE: removing federation flags.\n"
#	end_it 1
#}
#
#set fed_flags "None"
#array set clusters [get_clusterfed_info $fed_name]
#dict set clusters($fedc1) conn Self
#dict set clusters($fedc2) conn "Yes/Yes"
#dict set clusters($fedc3) conn "Yes/Yes"
#test_fed_status $fedc1 $fed_flags [array get clusters]


send_user "\n\nTest setting of cluster features\n"
set matches 0
set my_pid [spawn $sacctmgr -i modify cluster $fedc3 set features=fb,fa]
expect {
	-re "Setting$eol" {
		incr matches
		exp_continue
	}
	-re "\\s+Feature\\s+=\\s+fb$eol" {
		incr matches
		exp_continue
	}
	-re "\\s+Feature\\s+=\\s+fa$eol" {
		incr matches
		exp_continue
	}
	-re "Modified cluster...$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+$fedc3$eol" {
		incr matches
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sacctmgr add not responding\n"
		slow_kill $my_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$exit_code || $matches != 5} {
	send_user "$matches FAILURE: unexpected error.\n"
	end_it 1
}

array set clusters [get_clusterfed_info $fed_name]
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
dict set clusters($fedc3) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]

send_user "\n\nTest modifying/add of cluster features\n"
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$eol" {
		incr matches
		exp_continue
	}
	-re "Modified cluster...$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+$fedc3$eol" {
		incr matches
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sacctmgr add not responding\n"
		slow_kill $my_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$exit_code || $matches != 4} {
	send_user "$matches FAILURE: unexpected error.\n"
	end_it 1
}

array set clusters [get_clusterfed_info $fed_name]
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
dict set clusters($fedc3) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]

send_user "\n\nTest modifying/remove of cluster features\n"
set matches 0
set my_pid [spawn $sacctmgr -i modify cluster $fedc3 set features-=fb]
expect {
	-re "Setting$eol" {
		incr matches
		exp_continue
	}
	-re "\\s+Feature\\s+-=\\s+fb$eol" {
		incr matches
		exp_continue
	}
	-re "Modified cluster...$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+$fedc3$eol" {
		incr matches
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sacctmgr add not responding\n"
		slow_kill $my_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$exit_code || $matches != 4} {
	send_user "$matches FAILURE: unexpected error.\n"
	end_it 1
}

array set clusters [get_clusterfed_info $fed_name]
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
dict set clusters($fedc3) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]


send_user "\n\nTest clearing of cluster features\n"
set matches 0
set my_pid [spawn $sacctmgr -i modify cluster $fedc3 set features=]
expect {
	-re "Setting$eol" {
		incr matches
		exp_continue
	}
	-re "\\s+Feature\\s+=\\s+$eol" {
		incr matches
		exp_continue
	}
	-re "Modified cluster...$eol" {
		incr matches
		exp_continue
	}
	-re "^\\s+$fedc3$eol" {
		incr matches
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: sacctmgr add not responding\n"
		slow_kill $my_pid
		set exit_code 1
	}
	eof {
		wait
	}
}
if {$exit_code || $matches != 4} {
	send_user "$matches FAILURE: unexpected error.\n"
	end_it 1
}

array set clusters [get_clusterfed_info $fed_name]
dict set clusters($fedc1) conn Self
dict set clusters($fedc2) conn "Yes/Yes"
dict set clusters($fedc3) conn "Yes/Yes"
test_fed_status $fedc1 $fed_flags [array get clusters]

send_user "\n\ndelete federation and test from $fedc1\n"
delete_federations $fed_name
array unset clusters
test_fed_status $fedc2 $fed_flags [array get clusters]




end_it 0
