#! /usr/bin/env perl

##**************************************************************
##
## Copyright (C) 1990-2007, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
## 
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
## 
##    http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************

use strict;
use warnings;

use CondorTest;
use CondorUtils;
use Check::SimpleJob;

my $testname = 'job_router_xform';

# this test submits jobs and verifies that they ran where they were expected to.
# Some of the jobs will run directly, and some will be managed and routed.
# The router configuration has multiple routes, some of which should never be used
# The startd configuration has multiple slot types to simulate multiple execution pools

#testreq: personal
my $config = <<'CONDOR_TESTREQ_CONFIG';

DAEMON_LIST = $(DAEMON_LIST) JOB_ROUTER
#uncomment these to valgrind the job router (requires the 8.9.7 patch where daemons default to foreground mode)
#JOB_ROUTER = /usr/bin/valgrind
#JOB_ROUTER_ARGS = --leak-check=full --log-file=$(LOG)/job_router-vg.%p --error-limit=no $(LIBEXEC)/condor_job_router -f

NUM_CPUS = 5

# (some) jobs will be routed to production or analysis resources, some jobs will be left unrouted
# create 3 types of slots:  production, analysis, and unrouted
STARTD_ATTRS = PoolName
PoolName = "Hare"

NUM_SLOTS_TYPE_1 = 2
SLOT_TYPE_1_PARTITIONABLE = false
SLOT_TYPE_1 = 20%
SLOT_TYPE_1_START = target.ProductionJob
SLOT_TYPE_1_NAME_PREFIX = Prod

NUM_SLOTS_TYPE_2 = 2
SLOT_TYPE_2_PARTITIONABLE = false
SLOT_TYPE_2 = 20%
SLOT_TYPE_2_START = ! target.ProductionJob
SLOT_TYPE_2_NAME_PREFIX = Ana

NUM_SLOTS_TYPE_3 = 1
SLOT_TYPE_3_PARTITIONABLE = false
SLOT_TYPE_3 = 20%
SLOT_TYPE_3_START = ! target.RouteMe
SLOT_TYPE_3_NAME_PREFIX = Slow
SLOT_TYPE_3_PoolName = "Turtle"

JOB_ROUTER_POLLING_PERIOD = 10

JOB_ROUTER_PRE_ROUTE_TRANSFORM_NAMES = Defaults
JOB_ROUTER_TRANSFORM_Defaults @=jrd
  DELETE RouteMe
  SET RoutedJob true
  SET Requirements true
  SET CollHost "$(COLLECTOR_NAME)"
@jrd

# regression test for jira902
COLLECTOR_NAME = jira902

# Route even procid's to production, odd procid's to analysis
# Production1 and Analysis3 jobs require odd slot ids, Production2 and Analysis4 jobs require even slot ids
#
JOB_ROUTER_ROUTE_NAMES = Production1 Analysis3 Production2 Analysis4

JOB_ROUTER_ROUTE_Production1 @=rt
  UNIVERSE Vanilla
  REQUIREMENTS RouteMe && (ProcId & 1)
  #future: set MaxJobs and alter the ranat check accordingly
  #MaxJobs = 1
  SET ProductionJob 1
  SET ProductionRun 1
  SET Requirements target.slotId & 1
@rt

JOB_ROUTER_ROUTE_Production2 @=rt
  UNIVERSE VANILLA
  REQUIREMENTS RouteMe && (ProcId & 1)
  SET ProductionJob 1
  SET ProductionRun 2
  SET Requirements !(target.slotId & 1)
@rt

JOB_ROUTER_ROUTE_Analysis3 @=rt
  UNIVERSE 5
  REQUIREMENTS RouteMe && !(ProcId & 1)
  #future: set MaxJobs and alter the ranat check accordingly
  #MaxJobs = 1
  SET ProductionJob 0
  SET AnalysisRun 3
  SET Requirements target.slotId & 1
@rt

JOB_ROUTER_ROUTE_Analysis4 @=rt
  UNIVERSE vanilla
  Requirements RouteMe && !(ProcId & 1)
  SET ProductionJob 0
  SET AnalysisRun  4
  SET Requirements !(target.slotId & 1)
@rt

UPDATE_INTERVAL = 2
STARTER_UPDATE_INTERVAL = 2
SHADOW_QUEUE_UPDATE_INTERVAL = 2

CONDOR_TESTREQ_CONFIG
#endtestreq

my @rtconfig = (); # lines of job router config
CondorTest::runCondorTool("condor_job_router_info -config", \@rtconfig);

# We submit a cluster of 16 vanilla jobs
# half only run on the Turtle pool and are unrouted, half only run on the Hare pool
# the Hare pool only accepts jobs that have been routed to production or analysis
# The turtle pool only accepts jobs that request to be unrouted.  If all jobs run
# then we know that the job router did its job.
#
my $append_submit_commands = '
  My.RouteMe = $(route_me)
  My.Totem = $CHOICE(route_me, "Turtle", "Hare")
  My.RanAt = splitslotname(LastRemoteHost?:"_@")[0]
  Requirements = Target.PoolName == My.Totem
';

my $queue_args = 'route_me in (0, 1)';

if(is_windows()) {
	SimpleJob::RunCheck(
		runthis=>"sleep.exe",
		queue_sz => 8,
		multi_queue => $queue_args,
		append_submit_commands=>$append_submit_commands,
		);
} else {
	SimpleJob::RunCheck(
		queue_sz => 8,
		multi_queue => $queue_args,
		append_submit_commands=>$append_submit_commands,
		);
}

# give time for the history to update
sleep(10);

# at the end we will verify that each job went to the correct slot
# this verifies both correct routing, and correct non-routing
#
my %job = ();  # map of job info by id

my @jobs = (); # lines of output
my @historyjobs = (); # lines of output
CondorTest::runCondorTool("condor_q -af:j RouteMe RoutedJob Managed Totem RanAt ProductionJob ProductionRun AnalysisRun RoutedToJobId RoutedFromJobId CollHost", \@jobs);

# build a map of jobid to completion info
foreach my $line (@jobs) {
	fullchomp($line);
	my ($jobid,$info) = split / /, $line, 2;
	$job{$jobid} = $info;
	#print "$jobid = $info\n";
}

@jobs = ();
CondorTest::runCondorTool("condor_history -af:j RouteMe RoutedJob Managed Totem RanAt ProductionJob ProductionRun AnalysisRun RoutedToJobId RoutedFromJobId CollHost", \@jobs);

foreach my $line (@jobs) {
	fullchomp($line);
	my ($jobid,$info) = split / /, $line, 2;
	$job{$jobid} = $info;
	#print "$jobid = $info\n";
}

#my @ajob = ();
#CondorTest::runCondorTool("condor_q -limit 1 -const RouteMe -long ", \@ajob);
#@ajob = ();
#CondorTest::runCondorTool("condor_q -limit 1 -const ProductionJob -long ", \@ajob);

# Table of expected values for RunAt, the key is the concatenation of RouteMe + RoutedJob + Managed + ProductionJob
my %should_RunAt = (
	"0000" => 'Slow5',   # jobs 1.0 thru 1.7, not routed, "Turtle" jobs
	"1010" => '_',       # jobs 1.8 thru 1.15, routed, "Hare" jobs which run indirectly
	"0100" => 'Ana3',    # routed jobs of 1.8, 1.10, 1.12, 1.14
	"0101" => 'Prod1',   # routed jobs of 1.9, 1.11, 1.13, 1.15
	# these are keys that indicate that the pre-route transform did not run
	"1000" => 'NoPreRoute', # analysis job, but with RouteMe && ! RoutedJob
	"1001" => 'NoPreRoute', # production job, but with RouteMe && ! RoutedJob
);

# validate the map of jobs and where they ran
print "\n----- Checking final table of jobs and where they ran ----\n";
my ($jid, $info);
my $correct = 0;
my $incorrect = 0;
while (($jid, $info) = each(%job)) {
	my ($rm, $rj, $mg, $ttm, $ranat, $pr, $tags) = split / /, $info, 7;
	if ($rm eq "undefined") { $rm = "0"; }
	if ($rj eq "true") { $rj = 1; } else { $rj = 0; }
	if ($mg eq "undefined") { $mg = 0; } else { $mg = 1; }
	if ($pr eq "undefined") { $pr = 0; }
	my $key = $rm . $rj . $mg . $pr;
	print "$jid / $key / $ttm / $ranat / $pr $tags : ";
	# 
	my $should = "RoutingFailure";
	$should = $should_RunAt{$key} if (defined $should_RunAt{$key});
	if ($ranat eq $should) {
		$correct += 1;
		print "OK\n";
	} else {
		$incorrect += 1;
		print "FAILURE: should have run at $should\n";
	}
}

RegisterResult($correct == 24,test_name=>$testname,check_name => "All_jobs_ran_where_they_should");
RegisterResult($incorrect == 0,test_name=>$testname, check_name => "No_jobs_ran_where_they_should_not");

print "\n----- Checking CollHost attribute (jira902 regression) ----\n";

# jira902 regression test - check to see that CollHost is the expected value
@jobs = ();
CondorTest::runCondorTool("condor_q -const RoutedJob -af:j CollHost", \@jobs);

# It may be the case that the routed jobs have left the queue. If so, find them
# in the history file
@historyjobs = ();
if ($#jobs < 7) {
	CondorTest::runCondorTool("condor_history -const RoutedJob -af:j CollHost", \@historyjobs);
}

@jobs = (@jobs, @historyjobs);

$correct = 0;
$incorrect = 0;
foreach my $line (@jobs) {
	fullchomp($line);
	my ($jobid,$host) = split / /, $line, 2;
	if ($host eq 'jira902') {
		$correct += 1;
	} else {
		$incorrect += 1;
		print "\tFAILURE: $jobid CollHost=$host  should be jira902\n"
	}
}

RegisterResult($correct == 8,test_name=>$testname,check_name => "All_Correct_CollHost");
RegisterResult($incorrect == 0,test_name=>$testname, check_name => "No_Incorrect_CollHost");

CondorTest::EndTest();
