File indexing completed on 2026-08-12 08:24:58
0001 """Unit tests for ExecutionEngine module.
0002
0003 Tests cover:
0004 - JobContext (XCom push/pull, artifacts, logs)
0005 - Template (template substitution)
0006 - BashExecutionEngine (command execution, templating)
0007 - PythonExecutionEngine (function execution, arguments)
0008 - ContainerExecutionEngine (docker command building)
0009
0010 Project: AID2E v0.0.0
0011 """
0012
0013 import pytest
0014 from pathlib import Path
0015 import json
0016 import tempfile
0017 import os
0018
0019 from aid2e.utilities.workflows.execution_engine import (
0020 JobContext,
0021 StageContext,
0022 BranchContext,
0023 WorkflowSharedContext,
0024 BashExecutionEngine,
0025 PythonExecutionEngine,
0026 ContainerExecutionEngine,
0027 StackExecutionEngine,
0028 BaseExecutionEngine,
0029 Template,
0030 )
0031 from aid2e.utilities.configurations.experimental_stack_config import StackLayerConfig
0032
0033
0034 class TestJobContext:
0035 """Test JobContext (XCom and artifact management)."""
0036
0037 def test_task_context_init(self):
0038 """Test JobContext initialization."""
0039 context = JobContext(
0040 task_id='task_1',
0041 job_id='task_1',
0042 stage_id='stage_eval',
0043 workflow_id='workflow_1',
0044 design_point={'x': 0.5},
0045 execution_dir='/tmp'
0046 )
0047
0048 assert context.task_id == 'task_1'
0049 assert context.job_id == 'task_1'
0050 assert context.stage_id == 'stage_eval'
0051 assert context.workflow_id == 'workflow_1'
0052 assert context.design_point == {'x': 0.5}
0053 assert context.execution_dir == '/tmp'
0054 assert len(context.xcom) == 0
0055 assert len(context.artifacts) == 0
0056 assert len(context.logs) == 0
0057
0058 def test_xcom_push_and_pull(self):
0059 """Test XCom push/pull functionality."""
0060 context = JobContext(
0061 task_id='stage:upstream_task',
0062 job_id='upstream_task',
0063 stage_id='stage',
0064 workflow_id='workflow'
0065 )
0066
0067
0068 context.xcom_push('metrics', {'f1': 0.5, 'f2': 0.3})
0069 context.xcom_push('status', 'success')
0070
0071
0072 assert context.xcom_pull('stage:upstream_task', key='metrics') == {'f1': 0.5, 'f2': 0.3}
0073 assert context.xcom_pull('stage:upstream_task', key='status') == 'success'
0074 assert context.xcom_pull('stage:upstream_task', key='nonexistent') is None
0075
0076 def test_xcom_push_return_value_default(self):
0077 """Test default 'return_value' key in XCom."""
0078 context = JobContext(task_id='task_1', job_id='task_1', stage_id='s', workflow_id='w')
0079
0080
0081 context.xcom_push('return_value', 42)
0082
0083
0084 assert context.xcom_pull('task_1') == 42
0085
0086 def test_add_log(self):
0087 """Test log addition."""
0088 context = JobContext(task_id='s:t', job_id='t', stage_id='s', workflow_id='w')
0089
0090 context.add_log('Starting execution')
0091 context.add_log('Step 1 complete')
0092
0093 assert len(context.logs) == 2
0094 assert context.logs[0] == 'Starting execution'
0095 assert context.logs[1] == 'Step 1 complete'
0096
0097 def test_save_artifact(self):
0098 """Test artifact registration."""
0099 context = JobContext(task_id='s:t', job_id='t', stage_id='s', workflow_id='w')
0100
0101 context.save_artifact('objectives', '/work/objectives.json')
0102 context.save_artifact('metrics', '/work/metrics.json')
0103
0104 assert context.artifacts['objectives'] == '/work/objectives.json'
0105 assert context.artifacts['metrics'] == '/work/metrics.json'
0106
0107
0108 class TestTemplateSubstitutions:
0109 """Test Template."""
0110
0111 def test_substitutions(self):
0112 """Test common substitutions"""
0113 xcom = {
0114 'upstream:job:metric': 9,
0115 'upstream:job:sim:inputs': ['in_0.root', 'in_1.root', 'in_2.root'],
0116 'upstream:job:sim:outputs': ['out_0.root', 'out_1.root', 'out_2.root'],
0117 'upstream:job:sim:arguments': ['--crossingAngleBoost 0.025'],
0118 'upstream:job:return_value': {
0119 'stdout': 'Hello!',
0120 'stderr': '',
0121 'returncode': 9,
0122 },
0123 }
0124 artifacts = {
0125 'objective': '/output/here/objective.json',
0126 }
0127 workflow_context = WorkflowSharedContext(
0128 workflow_id='workflow',
0129 parameters = {'prepared_geometry_dir': '/geo/here'},
0130 )
0131 branch_context = BranchContext(
0132 branch_id='branch',
0133 parameters={},
0134 )
0135 stage_context = StageContext(
0136 stage_id='stage',
0137 parameters={},
0138 branch_context=branch_context,
0139 )
0140 job_context = JobContext(
0141 task_id=f'{stage_context.stage_id}:job',
0142 job_id='job',
0143 stage_id=stage_context.stage_id,
0144 workflow_id=workflow_context.workflow_id,
0145 design_point={'param_a': 'red', 'param_b': 'blue', 'param_c': 'green'},
0146 xcom=xcom,
0147 artifacts=artifacts,
0148 execution_dir='/execute/here',
0149 output_dir='/output/here',
0150 stage_context=stage_context,
0151 workflow_context=workflow_context,
0152 )
0153
0154 test_0 = "{{output_dir}}/out_{{design_point.param_a}}_{{design_point.param_b}}.root"
0155 test_1 = "{{execution_dir}}/{{branch_id}}_{{stage_id}}_{{job_id}}.log"
0156 test_2 = "{{geometry_dir}}/install/share/epic_{{workflow_id}}.xml"
0157 test_3 = "{{artifacts[objective]}}"
0158 test_4 = "{{xcom[upstream:job:metric]}}"
0159 test_5 = "{{xcom[upstream:job:sim:inputs](1)}}"
0160 test_6 = "{{xcom[upstream:job:return_value]('stdout')}}"
0161 test_7 = "{{inputs[upstream:job:sim](2)}}"
0162 test_8 = "{{outputs[upstream:job:sim](0)}}"
0163 test_9 = "{{arguments[upstream:job:sim](0)}}"
0164 assert Template.substitute(test_0, job_context) == "/output/here/out_red_blue.root"
0165 assert Template.substitute(test_1, job_context) == "/execute/here/branch_stage_job.log"
0166 assert Template.substitute(test_2, job_context) == "/geo/here/install/share/epic_workflow.xml"
0167 assert Template.substitute(test_3, job_context) == "/output/here/objective.json"
0168 assert Template.substitute(test_4, job_context) == "9"
0169 assert Template.substitute(test_5, job_context) == "in_1.root"
0170 assert Template.substitute(test_6, job_context) == "Hello!"
0171 assert Template.substitute(test_7, job_context) == "in_2.root"
0172 assert Template.substitute(test_8, job_context) == "out_0.root"
0173 assert Template.substitute(test_9, job_context) == "--crossingAngleBoost 0.025"
0174
0175
0176 class TestBashExecutionEngine:
0177 """Test BashExecutionEngine."""
0178
0179 def test_bash_engine_init(self):
0180 """Test BashExecutionEngine initialization."""
0181 op = BashExecutionEngine(
0182 engine_id='run_sim',
0183 bash_command='python script.py',
0184 env={'DEBUG': '1'}
0185 )
0186
0187 assert op.engine_id == 'run_sim'
0188 assert op.bash_command == 'python script.py'
0189 assert op.env == {'DEBUG': '1'}
0190
0191 def test_bash_engine_simple_command(self):
0192 """Test executing simple bash command."""
0193 context = JobContext(
0194 task_id='stage:test_bash',
0195 job_id='test_bash',
0196 stage_id='stage',
0197 workflow_id='workflow',
0198 execution_dir=None
0199 )
0200
0201 op = BashExecutionEngine(
0202 engine_id='echo_test',
0203 bash_command='echo "Hello World"'
0204 )
0205
0206 result = op.execute(context)
0207
0208 assert result['returncode'] == 0
0209 assert 'Hello World' in result['stdout']
0210
0211 def test_bash_engine_with_environment(self):
0212 """Test bash engine with environment variables."""
0213 context = JobContext(
0214 task_id='stage:test_env',
0215 job_id='test_env',
0216 stage_id='stage',
0217 workflow_id='workflow',
0218 execution_dir=None
0219 )
0220
0221 op = BashExecutionEngine(
0222 engine_id='check_env',
0223 bash_command='echo $TEST_VAR',
0224 env={'TEST_VAR': 'test_value'}
0225 )
0226
0227 result = op.execute(context)
0228 assert result['returncode'] == 0
0229
0230 def test_bash_engine_failure(self):
0231 """Test bash engine with failing command."""
0232 context = JobContext(
0233 task_id='stage:test_fail',
0234 job_id='test_fail',
0235 stage_id='stage',
0236 workflow_id='workflow',
0237 execution_dir=None
0238 )
0239
0240 op = BashExecutionEngine(
0241 engine_id='fail_test',
0242 bash_command='exit 1'
0243 )
0244
0245 with pytest.raises(RuntimeError, match='Command failed'):
0246 op.execute(context)
0247
0248 def test_bash_engine_template_substitution(self):
0249 """Test template variable substitution."""
0250 context = JobContext(
0251 task_id='stage:test_template',
0252 job_id='test_template',
0253 stage_id='stage',
0254 workflow_id='workflow',
0255 design_point={'input_file': 'test.json', 'output_file': 'out.json'}
0256 )
0257
0258 op = BashExecutionEngine(
0259 engine_id='template_test',
0260 bash_command='echo {{design_point.input_file}} {{design_point.output_file}}'
0261 )
0262
0263
0264 substituted = op._template.substitute(
0265 'Processing {{design_point.input_file}} to {{design_point.output_file}}',
0266 context
0267 )
0268
0269 assert substituted == 'Processing test.json to out.json'
0270
0271
0272 class TestPythonExecutionEngine:
0273 """Test PythonExecutionEngine."""
0274
0275 def test_python_engine_init(self):
0276 """Test PythonExecutionEngine initialization."""
0277 def my_func(context):
0278 return 42
0279
0280 op = PythonExecutionEngine(
0281 engine_id='compute',
0282 python_callable=my_func,
0283 op_kwargs={'param': 'value'}
0284 )
0285
0286 assert op.engine_id == 'compute'
0287 assert op.python_callable == my_func
0288 assert op.op_kwargs == {'param': 'value'}
0289
0290 def test_python_engine_simple_function(self):
0291 """Test executing simple Python function."""
0292 def compute_sum(context, a=1, b=2):
0293 return a + b
0294
0295 context = JobContext(
0296 task_id='stage:test_python',
0297 job_id='test_python',
0298 stage_id='stage',
0299 workflow_id='workflow'
0300 )
0301
0302 op = PythonExecutionEngine(
0303 engine_id='sum_test',
0304 python_callable=compute_sum,
0305 op_kwargs={'a': 5, 'b': 3}
0306 )
0307
0308 result = op.execute(context)
0309
0310 assert result == 8
0311 assert context.xcom_pull('stage:test_python', key='return_value') == 8
0312
0313 def test_python_engine_with_context(self):
0314 """Test Python function that uses JobContext."""
0315 def process_design_point(context):
0316 x = context.design_point.get('x', 0.0)
0317 y = context.design_point.get('y', 0.0)
0318 context.xcom_push('processed', {'sum': x + y})
0319 return {'sum': x + y, 'product': x * y}
0320
0321 context = JobContext(
0322 task_id='stage:processor',
0323 job_id='processor',
0324 stage_id='stage',
0325 workflow_id='workflow',
0326 design_point={'x': 3.0, 'y': 4.0}
0327 )
0328
0329 op = PythonExecutionEngine(
0330 engine_id='process',
0331 python_callable=process_design_point
0332 )
0333
0334 result = op.execute(context)
0335
0336 assert result['sum'] == 7.0
0337 assert result['product'] == 12.0
0338 assert context.xcom_pull('stage:processor', key='processed') == {'sum': 7.0}
0339
0340 def test_python_engine_exception(self):
0341 """Test Python engine with exception in function."""
0342 def failing_func(context):
0343 raise ValueError("Test error")
0344
0345 context = JobContext(
0346 task_id='stage:test_fail',
0347 job_id='test_fail',
0348 stage_id='stage',
0349 workflow_id='workflow'
0350 )
0351
0352 op = PythonExecutionEngine(
0353 engine_id='fail_test',
0354 python_callable=failing_func
0355 )
0356
0357 with pytest.raises(ValueError, match='Test error'):
0358 op.execute(context)
0359
0360 def test_python_engine_with_args_kwargs(self):
0361 """Test Python engine with positional and keyword arguments."""
0362 def multi_arg_func(context, a, b, c=10):
0363 return a + b + c
0364
0365 context = JobContext(
0366 task_id='stage:test_args',
0367 job_id='test_args',
0368 stage_id='stage',
0369 workflow_id='workflow'
0370 )
0371
0372 op = PythonExecutionEngine(
0373 engine_id='args_test',
0374 python_callable=multi_arg_func,
0375 op_args=(2, 3),
0376 op_kwargs={'c': 5}
0377 )
0378
0379 result = op.execute(context)
0380
0381 assert result == 10
0382
0383
0384 class TestContainerExecutionEngine:
0385 """Test ContainerExecutionEngine."""
0386
0387 def test_container_engine_init(self):
0388 """Test ContainerExecutionEngine initialization."""
0389 op = ContainerExecutionEngine(
0390 engine_id='run_container',
0391 image='myimage:1.0',
0392 command=['/app/run.sh'],
0393 environment={'VAR': 'value'},
0394 volumes={'/host': '/container'},
0395 resources={'memory': '4g'}
0396 )
0397
0398 assert op.engine_id == 'run_container'
0399 assert op.image == 'myimage:1.0'
0400 assert op.command == ['/app/run.sh']
0401 assert op.environment == {'VAR': 'value'}
0402 assert op.volumes == {'/host': '/container'}
0403 assert op.resources == {'memory': '4g'}
0404
0405 def test_container_engine_docker_command_basic(self):
0406 """Test basic docker command generation."""
0407 context = JobContext(
0408 task_id='stage:test',
0409 job_id='test',
0410 stage_id='stage',
0411 workflow_id='workflow',
0412 design_point={'x': 0.5}
0413 )
0414
0415 op = ContainerExecutionEngine(
0416 engine_id='docker_test',
0417 image='myimage:latest'
0418 )
0419
0420 cmd = op._build_docker_command(context)
0421
0422 assert 'docker run' in cmd
0423 assert '--rm' in cmd
0424 assert 'myimage:latest' in cmd
0425
0426 def test_container_engine_docker_with_env(self):
0427 """Test docker command with environment variables."""
0428 context = JobContext(
0429 task_id='stage:test',
0430 job_id='test',
0431 stage_id='stage',
0432 workflow_id='workflow',
0433 design_point={'x': 0.5}
0434 )
0435
0436 op = ContainerExecutionEngine(
0437 engine_id='docker_env',
0438 image='myimage:latest',
0439 environment={
0440 'PARAM1': 'value1',
0441 'PARAM2': '{design_point.x}'
0442 }
0443 )
0444
0445 cmd = op._build_docker_command(context)
0446
0447 assert '-e PARAM1=value1' in cmd
0448
0449 def test_container_engine_docker_with_volumes(self):
0450 """Test docker command with volume mounts."""
0451 context = JobContext(
0452 task_id='stage:test',
0453 job_id='test',
0454 stage_id='stage',
0455 workflow_id='workflow'
0456 )
0457
0458 op = ContainerExecutionEngine(
0459 engine_id='docker_vol',
0460 image='myimage:latest',
0461 volumes={
0462 '/host/data': '/data',
0463 '/host/output': '/output'
0464 }
0465 )
0466
0467 cmd = op._build_docker_command(context)
0468
0469 assert '-v /host/data:/data' in cmd
0470 assert '-v /host/output:/output' in cmd
0471
0472 def test_container_engine_docker_with_resources(self):
0473 """Test docker command with resource constraints."""
0474 context = JobContext(
0475 task_id='stage:test',
0476 job_id='test',
0477 stage_id='stage',
0478 workflow_id='workflow'
0479 )
0480
0481 op = ContainerExecutionEngine(
0482 engine_id='docker_res',
0483 image='myimage:latest',
0484 resources={
0485 'memory': '4g',
0486 'cpus': '2'
0487 }
0488 )
0489
0490 cmd = op._build_docker_command(context)
0491
0492 assert '-m 4g' in cmd
0493 assert '--cpus 2' in cmd
0494
0495 def test_container_engine_docker_with_command(self):
0496 """Test docker command with command override."""
0497 context = JobContext(
0498 task_id='stage:test',
0499 job_id='test',
0500 stage_id='stage',
0501 workflow_id='workflow'
0502 )
0503
0504 op = ContainerExecutionEngine(
0505 engine_id='docker_cmd',
0506 image='myimage:latest',
0507 command=['/app/script.sh', 'arg1', 'arg2']
0508 )
0509
0510 cmd = op._build_docker_command(context)
0511
0512 assert '/app/script.sh' in cmd
0513 assert 'arg1' in cmd
0514 assert 'arg2' in cmd
0515
0516
0517 class TestExecutionEngineInheritance:
0518 """Test execution engine inheritance and polymorphism."""
0519
0520 def test_all_engines_inherit_from_base(self):
0521 """Test that all engines inherit from BaseExecutionEngine."""
0522 assert issubclass(BashExecutionEngine, BaseExecutionEngine)
0523 assert issubclass(PythonExecutionEngine, BaseExecutionEngine)
0524 assert issubclass(ContainerExecutionEngine, BaseExecutionEngine)
0525 assert issubclass(StackExecutionEngine, BaseExecutionEngine)
0526
0527 def test_base_engine_repr(self):
0528 """Test string representation."""
0529 def dummy(context):
0530 pass
0531
0532 bash_op = BashExecutionEngine(engine_id='bash_task', bash_command='echo hi')
0533 python_op = PythonExecutionEngine(engine_id='python_task', python_callable=dummy)
0534 container_op = ContainerExecutionEngine(engine_id='container_task', image='img:1.0')
0535 stack_op = StackExecutionEngine(
0536 engine_id='stack_task',
0537 stack_type='epic',
0538 layers=[StackLayerConfig(name='sim_0', layer='sim', inputs=["in.root"], outputs=["out.root"])]
0539 )
0540
0541 assert 'BashExecutionEngine' in repr(bash_op)
0542 assert 'bash_task' in repr(bash_op)
0543 assert 'PythonExecutionEngine' in repr(python_op)
0544 assert 'ContainerExecutionEngine' in repr(container_op)
0545 assert 'StackExecutionEngine' in repr(stack_op)
0546
0547
0548 if __name__ == '__main__':
0549 pytest.main([__file__, '-v'])
