class TestMethodTrigger(unittest.TestCase):
def setUp(self):
# Create a basic object
class TestClass(object):
def func(self):
pass
self.testobj = TestClass()
# Install trigger
self.trigger = MethodTrigger(self.testobj, "func")
def test_no_call(self):
""" Do not trigger. """
# Wait for trigger
self.assertFalse(self.trigger.wait(0.1))
def test_serial_call(self):
""" Trigger and wait serially. """
# Call function
self.testobj.func()
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
def test_parallel_call(self):
""" Trigger and wait in parallel. """
class wait_thread(Thread):
def __init__(self, trigger):
Thread.__init__(self)
self.result = None
self.trigger = trigger
def run(self):
self.result = self.trigger.wait(1.0)
class call_thread(Thread):
def __init__(self,func):
Thread.__init__(self)
self.func = func
def run(self):
time.sleep(0.5)
self.func()
# Start threads
t1 = wait_thread(self.trigger)
t1.start()
t2 = call_thread(self.testobj.func)
t2.start()
# Wait for them to finish
t1.join()
t2.join()
# Inspect result
self.assertTrue(t1.result)
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, {
"create": "always",
"delete": "always",
"node": {
"type": "topic"
}
})
self.bus.open()
self.addCleanup(self.bus.close)
# Define the services we use
self.status_service = "%s/%s.TaskGetSpecification" % (
self.busname, DEFAULT_OTDB_SERVICENAME)
# ================================
# Setup mock parset service
# ================================
# Nr of parsets requested, to detect multiple requests for the same parset, or of superfluous parsets
self.requested_parsets = 0
def TaskSpecificationService(OtdbID):
if OtdbID == 1:
predecessors = "[2,3]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
self.requested_parsets += 1
return {
"Version.number": "1",
PARSET_PREFIX + "Observation.ObsID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors":
predecessors,
}
parset_service = Service(DEFAULT_OTDB_SERVICENAME +
".TaskSpecification",
TaskSpecificationService,
busname=self.busname)
parset_service.start_listening()
self.addCleanup(parset_service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = RATaskSpecifiedBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onTaskSpecified")
class TestMethodTrigger(unittest.TestCase):
def setUp(self):
# Create a basic object
class TestClass(object):
def func(self):
pass
self.testobj = TestClass()
# Install trigger
self.trigger = MethodTrigger(self.testobj, "func")
def test_no_call(self):
""" Do not trigger. """
# Wait for trigger
self.assertFalse(self.trigger.wait(0.1))
def test_serial_call(self):
""" Trigger and wait serially. """
# Call function
self.testobj.func()
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
def test_parallel_call(self):
""" Trigger and wait in parallel. """
class wait_thread(Thread):
def __init__(self, trigger):
Thread.__init__(self)
self.result = None
self.trigger = trigger
def run(self):
self.result = self.trigger.wait(1.0)
class call_thread(Thread):
def __init__(self, func):
Thread.__init__(self)
self.func = func
def run(self):
time.sleep(0.5)
self.func()
# Start threads
t1 = wait_thread(self.trigger)
t1.start()
t2 = call_thread(self.testobj.func)
t2.start()
# Wait for them to finish
t1.join()
t2.join()
# Inspect result
self.assertTrue(t1.result)
def setUp(self):
# Create a basic object
class TestClass(object):
def func(self):
pass
self.testobj = TestClass()
# Install trigger
self.trigger = MethodTrigger(self.testobj, "func")
class TestArgs(unittest.TestCase):
def setUp(self):
# Create a basic object
class TestClass(object):
def func(self, a, b, c=None, d=None):
pass
self.testobj = TestClass()
# Install trigger
self.trigger = MethodTrigger(self.testobj, "func")
def test_args(self):
""" Trigger and check args. """
# Call function
self.testobj.func(1, 2)
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
# Check stored arguments
self.assertEqual(self.trigger.args, (1, 2))
def test_kwargs(self):
""" Trigger and check kwargs. """
# Call function
self.testobj.func(a=1, b=2)
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
# Check stored arguments
self.assertEqual(self.trigger.kwargs, {"a": 1, "b": 2})
def test_full(self):
""" Trigger and check both args and kwargs. """
# Call function
self.testobj.func(1, 2, c=3, d=4)
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
# Check stored arguments
self.assertEqual(self.trigger.args, (1, 2))
self.assertEqual(self.trigger.kwargs, {"c": 3, "d": 4})
class TestArgs(unittest.TestCase):
def setUp(self):
# Create a basic object
class TestClass(object):
def func(self, a, b, c=None, d=None):
pass
self.testobj = TestClass()
# Install trigger
self.trigger = MethodTrigger(self.testobj, "func")
def test_args(self):
""" Trigger and check args. """
# Call function
self.testobj.func(1, 2)
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
# Check stored arguments
self.assertEqual(self.trigger.args, (1, 2))
def test_kwargs(self):
""" Trigger and check kwargs. """
# Call function
self.testobj.func(a=1, b=2)
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
# Check stored arguments
self.assertEqual(self.trigger.kwargs, {"a": 1, "b": 2})
def test_full(self):
""" Trigger and check both args and kwargs. """
# Call function
self.testobj.func(1, 2, c=3, d=4)
# Wait for trigger
self.assertTrue(self.trigger.wait(0.1))
# Check stored arguments
self.assertEqual(self.trigger.args, (1, 2))
self.assertEqual(self.trigger.kwargs, {"c": 3, "d": 4})
def setUp(self):
# Create a basic object
class TestClass(object):
def func(self):
pass
self.testobj = TestClass()
# Install trigger
self.trigger = MethodTrigger(self.testobj, "func")
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, { "create": "always", "delete": "always", "node": { "type": "topic" } })
self.bus.open()
self.addCleanup(self.bus.close)
# Define the services we use
self.status_service = "%s/%s.TaskGetSpecification" % (self.busname,DEFAULT_OTDB_SERVICENAME)
# ================================
# Setup mock parset service
# ================================
# Nr of parsets requested, to detect multiple requests for the same parset, or of superfluous parsets
self.requested_parsets = 0
def TaskSpecificationService( OtdbID ):
if OtdbID == 1:
predecessors = "[2,3]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
self.requested_parsets += 1
return {
"Version.number": "1",
PARSET_PREFIX + "Observation.ObsID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors": predecessors,
}
parset_service = Service(DEFAULT_OTDB_SERVICENAME+".TaskSpecification", TaskSpecificationService, busname=self.busname)
parset_service.start_listening()
self.addCleanup(parset_service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = RATaskSpecifiedBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onTaskSpecified")
class TestService(unittest.TestCase):
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, {
"create": "always",
"delete": "always",
"node": {
"type": "topic"
}
})
self.bus.open()
self.addCleanup(self.bus.close)
# Define the services we use
self.status_service = "%s/%s.TaskGetSpecification" % (
self.busname, DEFAULT_OTDB_SERVICENAME)
# ================================
# Setup mock parset service
# ================================
# Nr of parsets requested, to detect multiple requests for the same parset, or of superfluous parsets
self.requested_parsets = 0
def TaskSpecificationService(OtdbID):
if OtdbID == 1:
predecessors = "[2,3]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
self.requested_parsets += 1
return {
"Version.number": "1",
PARSET_PREFIX + "Observation.ObsID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors":
predecessors,
}
parset_service = Service(DEFAULT_OTDB_SERVICENAME +
".TaskSpecification",
TaskSpecificationService,
busname=self.busname)
parset_service.start_listening()
self.addCleanup(parset_service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = RATaskSpecifiedBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onTaskSpecified")
def testNoPredecessors(self):
"""
Request the resources for a simulated obsid 3, with the following predecessor tree:
3 requires nothing
"""
with RATaskSpecified(otdb_notification_busname=self.busname,
otdb_service_busname=self.busname,
notification_busname=self.busname) as jts:
# Send fake status update
with ToBus(self.status_service) as tb:
msg = EventMessage(
content={
"treeID": 3,
"state": "prescheduled",
"time_of_change": datetime.datetime(2016, 1, 1),
})
tb.send(msg)
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 3)
resourceIndicators = self.trigger.args[2]
self.assertNotIn("1", resourceIndicators)
self.assertNotIn("2", resourceIndicators)
self.assertIn("3", resourceIndicators)
# Make sure we only requested one parset
self.assertEqual(self.requested_parsets, 1)
def testPredecessors(self):
"""
Request the resources for a simulated obsid 1, with the following predecessor tree:
1 requires 2, 3
2 requires 3
3 requires nothing
"""
with RATaskSpecified(otdb_notification_busname=self.busname,
notification_busname=self.busname) as jts:
# Send fake status update
with ToBus(self.status_service) as tb:
msg = EventMessage(
content={
"treeID": 1,
"state": "prescheduled",
"time_of_change": datetime.datetime(2016, 1, 1),
})
tb.send(msg)
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 1)
resourceIndicators = self.trigger.args[2]
self.assertIn("1", resourceIndicators)
self.assertIn("2", resourceIndicators)
self.assertIn("3", resourceIndicators)
# Make sure we only requested exactly three parsets
self.assertEqual(self.requested_parsets, 3)
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, { "create": "always", "delete": "always", "node": { "type": "topic" } })
self.bus.open()
self.addCleanup(self.bus.close)
# Patch SLURM
class MockSlurm(object):
def __init__(self, *args, **kwargs):
self.scheduled_jobs = {}
def submit(self, jobName, *args, **kwargs):
print "Schedule SLURM job '%s': %s, %s" % (jobName, args, kwargs)
self.scheduled_jobs[jobName] = (args, kwargs)
# Return job ID
return "42"
def jobid(self, jobname):
if jobname in ["1", "2", "3"]:
return jobname
# "4" is an observation, so no SLURM job
return None
patcher = patch('lofar.mac.PipelineControl.Slurm')
patcher.start().side_effect = MockSlurm
self.addCleanup(patcher.stop)
# Catch functions to prevent running executables
patcher = patch('lofar.mac.PipelineControl.Parset.dockerImage')
patcher.start().return_value = "lofar-pipeline:trunk"
self.addCleanup(patcher.stop)
# ================================
# Setup mock otdb service
# ================================
class MockOTDBService(MessageHandlerInterface):
def __init__(self, notification_bus):
"""
notification_bus: bus to send state changes to
"""
super(MockOTDBService, self).__init__()
self.service2MethodMap = {
"TaskGetSpecification": self.TaskGetSpecification,
"TaskSetStatus": self.TaskSetStatus,
}
self.notification_bus = notification_bus
def TaskGetSpecification(self, OtdbID):
print "***** TaskGetSpecification(%s) *****" % (OtdbID,)
if OtdbID == 1:
predecessors = "[2,3,4]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
elif OtdbID == 4:
return { "TaskSpecification": {
"Version.number": "1",
PARSET_PREFIX + "Observation.otdbID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors": "[]",
PARSET_PREFIX + "Observation.processType": "Observation",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName": "CEP4",
PARSET_PREFIX + "Observation.stopTime": "2016-01-01 01:00:00",
} }
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
return { "TaskSpecification": {
"Version.number": "1",
PARSET_PREFIX + "Observation.otdbID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors": predecessors,
PARSET_PREFIX + "Observation.processType": "Pipeline",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName": "CEP4",
} }
def TaskSetStatus(self, OtdbID, NewStatus, UpdateTimestamps):
print "***** TaskSetStatus(%s,%s) *****" % (OtdbID, NewStatus)
# Broadcast the state change
content = { "treeID" : OtdbID, "state" : NewStatus, "time_of_change" : datetime.datetime.utcnow() }
msg = EventMessage(context=DEFAULT_OTDB_NOTIFICATION_SUBJECT, content=content)
self.notification_bus.send(msg)
return {'OtdbID':OtdbID, 'MomID':None, 'Success':True}
service = Service(DEFAULT_OTDB_SERVICENAME,
MockOTDBService,
busname=self.busname,
use_service_methods=True,
handler_args={ "notification_bus": self.bus })
service.start_listening()
self.addCleanup(service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = OTDBBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onObservationQueued")
class TestPipelineControl(unittest.TestCase):
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, { "create": "always", "delete": "always", "node": { "type": "topic" } })
self.bus.open()
self.addCleanup(self.bus.close)
# Patch SLURM
class MockSlurm(object):
def __init__(self, *args, **kwargs):
self.scheduled_jobs = {}
def submit(self, jobName, *args, **kwargs):
print "Schedule SLURM job '%s': %s, %s" % (jobName, args, kwargs)
self.scheduled_jobs[jobName] = (args, kwargs)
# Return job ID
return "42"
def jobid(self, jobname):
if jobname in ["1", "2", "3"]:
return jobname
# "4" is an observation, so no SLURM job
return None
patcher = patch('lofar.mac.PipelineControl.Slurm')
patcher.start().side_effect = MockSlurm
self.addCleanup(patcher.stop)
# Catch functions to prevent running executables
patcher = patch('lofar.mac.PipelineControl.Parset.dockerImage')
patcher.start().return_value = "lofar-pipeline:trunk"
self.addCleanup(patcher.stop)
# ================================
# Setup mock otdb service
# ================================
class MockOTDBService(MessageHandlerInterface):
def __init__(self, notification_bus):
"""
notification_bus: bus to send state changes to
"""
super(MockOTDBService, self).__init__()
self.service2MethodMap = {
"TaskGetSpecification": self.TaskGetSpecification,
"TaskSetStatus": self.TaskSetStatus,
}
self.notification_bus = notification_bus
def TaskGetSpecification(self, OtdbID):
print "***** TaskGetSpecification(%s) *****" % (OtdbID,)
if OtdbID == 1:
predecessors = "[2,3,4]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
elif OtdbID == 4:
return { "TaskSpecification": {
"Version.number": "1",
PARSET_PREFIX + "Observation.otdbID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors": "[]",
PARSET_PREFIX + "Observation.processType": "Observation",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName": "CEP4",
PARSET_PREFIX + "Observation.stopTime": "2016-01-01 01:00:00",
} }
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
return { "TaskSpecification": {
"Version.number": "1",
PARSET_PREFIX + "Observation.otdbID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors": predecessors,
PARSET_PREFIX + "Observation.processType": "Pipeline",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName": "CEP4",
} }
def TaskSetStatus(self, OtdbID, NewStatus, UpdateTimestamps):
print "***** TaskSetStatus(%s,%s) *****" % (OtdbID, NewStatus)
# Broadcast the state change
content = { "treeID" : OtdbID, "state" : NewStatus, "time_of_change" : datetime.datetime.utcnow() }
msg = EventMessage(context=DEFAULT_OTDB_NOTIFICATION_SUBJECT, content=content)
self.notification_bus.send(msg)
return {'OtdbID':OtdbID, 'MomID':None, 'Success':True}
service = Service(DEFAULT_OTDB_SERVICENAME,
MockOTDBService,
busname=self.busname,
use_service_methods=True,
handler_args={ "notification_bus": self.bus })
service.start_listening()
self.addCleanup(service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = OTDBBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onObservationQueued")
def test_setStatus(self):
with PipelineControl(otdb_notification_busname=self.busname, otdb_service_busname=self.busname) as pipelineControl:
pipelineControl._setStatus(12345, "queued")
# Wait for the status to propagate
self.assertTrue(self.trigger.wait())
self.assertEqual(self.trigger.args[0], 12345)
def testNoPredecessors(self):
"""
Request to start a simulated obsid 3, with the following predecessor tree:
3 requires nothing
"""
with PipelineControl(otdb_notification_busname=self.busname, otdb_service_busname=self.busname) as pipelineControl:
# Send fake status update
pipelineControl._setStatus(3, "scheduled")
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 3) # otdbId
# Check if job was scheduled
self.assertIn("3", pipelineControl.slurm.scheduled_jobs)
self.assertIn("3-aborted", pipelineControl.slurm.scheduled_jobs)
def testPredecessors(self):
"""
Request to start a simulated obsid 1, with the following predecessor tree:
1 requires 2, 3, 4
2 requires 3
4 is an observation
"""
with PipelineControl(otdb_notification_busname=self.busname, otdb_service_busname=self.busname) as pipelineControl:
# Send fake status update
pipelineControl._setStatus(1, "scheduled")
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 1) # otdbId
# Check if job was scheduled
self.assertIn("1", pipelineControl.slurm.scheduled_jobs)
self.assertIn("1-aborted", pipelineControl.slurm.scheduled_jobs)
# Earliest start of this job > stop time of observation
for p in pipelineControl.slurm.scheduled_jobs["1"][1]["sbatch_params"]:
if p.startswith("--begin="):
begin = datetime.datetime.strptime(p, "--begin=%Y-%m-%dT%H:%M:%S")
self.assertGreater(begin, datetime.datetime(2016, 1, 1, 1, 0, 0))
break
else:
self.assertTrue(False, "--begin parameter not given to SLURM job")
class TestService(unittest.TestCase):
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, { "create": "always", "delete": "always", "node": { "type": "topic" } })
self.bus.open()
self.addCleanup(self.bus.close)
# Define the services we use
self.status_service = "%s/%s.TaskGetSpecification" % (self.busname,DEFAULT_OTDB_SERVICENAME)
# ================================
# Setup mock parset service
# ================================
# Nr of parsets requested, to detect multiple requests for the same parset, or of superfluous parsets
self.requested_parsets = 0
def TaskSpecificationService( OtdbID ):
if OtdbID == 1:
predecessors = "[2,3]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
self.requested_parsets += 1
return {
"Version.number": "1",
PARSET_PREFIX + "Observation.ObsID": str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors": predecessors,
}
parset_service = Service(DEFAULT_OTDB_SERVICENAME+".TaskSpecification", TaskSpecificationService, busname=self.busname)
parset_service.start_listening()
self.addCleanup(parset_service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = RATaskSpecifiedBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onTaskSpecified")
def testNoPredecessors(self):
"""
Request the resources for a simulated obsid 3, with the following predecessor tree:
3 requires nothing
"""
with RATaskSpecified(otdb_notification_busname=self.busname,
otdb_service_busname=self.busname,
notification_busname=self.busname) as jts:
# Send fake status update
with ToBus(self.status_service) as tb:
msg = EventMessage(content={
"treeID": 3,
"state": "prescheduled",
"time_of_change": datetime.datetime(2016,1,1),
})
tb.send(msg)
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 3)
resourceIndicators = self.trigger.args[2]
self.assertNotIn("1", resourceIndicators)
self.assertNotIn("2", resourceIndicators)
self.assertIn("3", resourceIndicators);
# Make sure we only requested one parset
self.assertEqual(self.requested_parsets, 1)
def testPredecessors(self):
"""
Request the resources for a simulated obsid 1, with the following predecessor tree:
1 requires 2, 3
2 requires 3
3 requires nothing
"""
with RATaskSpecified(otdb_notification_busname=self.busname, notification_busname=self.busname) as jts:
# Send fake status update
with ToBus(self.status_service) as tb:
msg = EventMessage(content={
"treeID": 1,
"state": "prescheduled",
"time_of_change": datetime.datetime(2016,1,1),
})
tb.send(msg)
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 1)
resourceIndicators = self.trigger.args[2]
self.assertIn("1", resourceIndicators);
self.assertIn("2", resourceIndicators);
self.assertIn("3", resourceIndicators);
# Make sure we only requested exactly three parsets
self.assertEqual(self.requested_parsets, 3)
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, {
"create": "always",
"delete": "always",
"node": {
"type": "topic"
}
})
self.bus.open()
self.addCleanup(self.bus.close)
# Patch SLURM
class MockSlurm(object):
def __init__(self, *args, **kwargs):
self.scheduled_jobs = {}
def submit(self, jobName, *args, **kwargs):
print "Schedule SLURM job '%s': %s, %s" % (jobName, args,
kwargs)
self.scheduled_jobs[jobName] = (args, kwargs)
# Return job ID
return "42"
def jobid(self, jobname):
if jobname in ["1", "2", "3"]:
return jobname
# "4" is an observation, so no SLURM job
return None
patcher = patch('lofar.mac.PipelineControl.Slurm')
patcher.start().side_effect = MockSlurm
self.addCleanup(patcher.stop)
# Catch functions to prevent running executables
patcher = patch('lofar.mac.PipelineControl.Parset.dockerImage')
patcher.start().return_value = "lofar-pipeline:trunk"
self.addCleanup(patcher.stop)
# ================================
# Setup mock otdb service
# ================================
class MockOTDBService(MessageHandlerInterface):
def __init__(self, notification_bus):
"""
notification_bus: bus to send state changes to
"""
super(MockOTDBService, self).__init__()
self.service2MethodMap = {
"TaskGetSpecification": self.TaskGetSpecification,
"TaskSetStatus": self.TaskSetStatus,
}
self.notification_bus = notification_bus
def TaskGetSpecification(self, OtdbID):
print "***** TaskGetSpecification(%s) *****" % (OtdbID, )
if OtdbID == 1:
predecessors = "[2,3,4]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
elif OtdbID == 4:
return {
"TaskSpecification": {
"Version.number":
"1",
PARSET_PREFIX + "Observation.otdbID":
str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors":
"[]",
PARSET_PREFIX + "Observation.processType":
"Observation",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName":
"CEP4",
PARSET_PREFIX + "Observation.stopTime":
"2016-01-01 01:00:00",
}
}
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
return {
"TaskSpecification": {
"Version.number":
"1",
PARSET_PREFIX + "Observation.otdbID":
str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors":
predecessors,
PARSET_PREFIX + "Observation.processType":
"Pipeline",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName":
"CEP4",
}
}
def TaskSetStatus(self, OtdbID, NewStatus, UpdateTimestamps):
print "***** TaskSetStatus(%s,%s) *****" % (OtdbID, NewStatus)
# Broadcast the state change
content = {
"treeID": OtdbID,
"state": NewStatus,
"time_of_change": datetime.datetime.utcnow()
}
msg = EventMessage(context=DEFAULT_OTDB_NOTIFICATION_SUBJECT,
content=content)
self.notification_bus.send(msg)
return {'OtdbID': OtdbID, 'MomID': None, 'Success': True}
service = Service(DEFAULT_OTDB_SERVICENAME,
MockOTDBService,
busname=self.busname,
use_service_methods=True,
handler_args={"notification_bus": self.bus})
service.start_listening()
self.addCleanup(service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = OTDBBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onObservationQueued")
class TestPipelineControl(unittest.TestCase):
def setUp(self):
# Create a random bus
self.busname = "%s-%s" % (sys.argv[0], str(uuid.uuid4())[:8])
self.bus = ToBus(self.busname, {
"create": "always",
"delete": "always",
"node": {
"type": "topic"
}
})
self.bus.open()
self.addCleanup(self.bus.close)
# Patch SLURM
class MockSlurm(object):
def __init__(self, *args, **kwargs):
self.scheduled_jobs = {}
def submit(self, jobName, *args, **kwargs):
print "Schedule SLURM job '%s': %s, %s" % (jobName, args,
kwargs)
self.scheduled_jobs[jobName] = (args, kwargs)
# Return job ID
return "42"
def jobid(self, jobname):
if jobname in ["1", "2", "3"]:
return jobname
# "4" is an observation, so no SLURM job
return None
patcher = patch('lofar.mac.PipelineControl.Slurm')
patcher.start().side_effect = MockSlurm
self.addCleanup(patcher.stop)
# Catch functions to prevent running executables
patcher = patch('lofar.mac.PipelineControl.Parset.dockerImage')
patcher.start().return_value = "lofar-pipeline:trunk"
self.addCleanup(patcher.stop)
# ================================
# Setup mock otdb service
# ================================
class MockOTDBService(MessageHandlerInterface):
def __init__(self, notification_bus):
"""
notification_bus: bus to send state changes to
"""
super(MockOTDBService, self).__init__()
self.service2MethodMap = {
"TaskGetSpecification": self.TaskGetSpecification,
"TaskSetStatus": self.TaskSetStatus,
}
self.notification_bus = notification_bus
def TaskGetSpecification(self, OtdbID):
print "***** TaskGetSpecification(%s) *****" % (OtdbID, )
if OtdbID == 1:
predecessors = "[2,3,4]"
elif OtdbID == 2:
predecessors = "[3]"
elif OtdbID == 3:
predecessors = "[]"
elif OtdbID == 4:
return {
"TaskSpecification": {
"Version.number":
"1",
PARSET_PREFIX + "Observation.otdbID":
str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors":
"[]",
PARSET_PREFIX + "Observation.processType":
"Observation",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName":
"CEP4",
PARSET_PREFIX + "Observation.stopTime":
"2016-01-01 01:00:00",
}
}
else:
raise Exception("Invalid OtdbID: %s" % OtdbID)
return {
"TaskSpecification": {
"Version.number":
"1",
PARSET_PREFIX + "Observation.otdbID":
str(OtdbID),
PARSET_PREFIX + "Observation.Scheduler.predecessors":
predecessors,
PARSET_PREFIX + "Observation.processType":
"Pipeline",
PARSET_PREFIX + "Observation.Cluster.ProcessingCluster.clusterName":
"CEP4",
}
}
def TaskSetStatus(self, OtdbID, NewStatus, UpdateTimestamps):
print "***** TaskSetStatus(%s,%s) *****" % (OtdbID, NewStatus)
# Broadcast the state change
content = {
"treeID": OtdbID,
"state": NewStatus,
"time_of_change": datetime.datetime.utcnow()
}
msg = EventMessage(context=DEFAULT_OTDB_NOTIFICATION_SUBJECT,
content=content)
self.notification_bus.send(msg)
return {'OtdbID': OtdbID, 'MomID': None, 'Success': True}
service = Service(DEFAULT_OTDB_SERVICENAME,
MockOTDBService,
busname=self.busname,
use_service_methods=True,
handler_args={"notification_bus": self.bus})
service.start_listening()
self.addCleanup(service.stop_listening)
# ================================
# Setup listener to catch result
# of our service
# ================================
listener = OTDBBusListener(busname=self.busname)
listener.start_listening()
self.addCleanup(listener.stop_listening)
self.trigger = MethodTrigger(listener, "onObservationQueued")
def test_setStatus(self):
with PipelineControl(
otdb_notification_busname=self.busname,
otdb_service_busname=self.busname) as pipelineControl:
pipelineControl._setStatus(12345, "queued")
# Wait for the status to propagate
self.assertTrue(self.trigger.wait())
self.assertEqual(self.trigger.args[0], 12345)
def testNoPredecessors(self):
"""
Request to start a simulated obsid 3, with the following predecessor tree:
3 requires nothing
"""
with PipelineControl(
otdb_notification_busname=self.busname,
otdb_service_busname=self.busname) as pipelineControl:
# Send fake status update
pipelineControl._setStatus(3, "scheduled")
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 3) # otdbId
# Check if job was scheduled
self.assertIn("3", pipelineControl.slurm.scheduled_jobs)
self.assertIn("3-aborted", pipelineControl.slurm.scheduled_jobs)
def testPredecessors(self):
"""
Request to start a simulated obsid 1, with the following predecessor tree:
1 requires 2, 3, 4
2 requires 3
4 is an observation
"""
with PipelineControl(
otdb_notification_busname=self.busname,
otdb_service_busname=self.busname) as pipelineControl:
# Send fake status update
pipelineControl._setStatus(1, "scheduled")
# Wait for message to arrive
self.assertTrue(self.trigger.wait())
# Verify message
self.assertEqual(self.trigger.args[0], 1) # otdbId
# Check if job was scheduled
self.assertIn("1", pipelineControl.slurm.scheduled_jobs)
self.assertIn("1-aborted", pipelineControl.slurm.scheduled_jobs)
# Earliest start of this job > stop time of observation
for p in pipelineControl.slurm.scheduled_jobs["1"][1][
"sbatch_params"]:
if p.startswith("--begin="):
begin = datetime.datetime.strptime(
p, "--begin=%Y-%m-%dT%H:%M:%S")
self.assertGreater(begin,
datetime.datetime(2016, 1, 1, 1, 0, 0))
break
else:
self.assertTrue(False,
"--begin parameter not given to SLURM job")