Coverage for agentos/tests/test_workflow.py: 0%

1"""Comprehensive tests for agentos/workflow/__init__.py.""" 

2 

3import json 

4 

5import pytest 

6 

7from agentos.workflow import ( 

8 ConditionEvaluator, 

9 ConditionOperator, 

10 ErrorStrategy, 

11 ExecutionStatus, 

12 StepResult, 

13 StepType, 

14 WorkflowContext, 

15 WorkflowDefinition, 

16 WorkflowEngine, 

17 WorkflowParser, 

18 WorkflowStep, 

19 WorkflowTemplates, 

20) 

21 

22# ============================================================================ 

23# Enums 

24# ============================================================================ 

25 

26class TestStepType: 

27 def test_values(self): 

28 assert StepType.TASK.value == "task" 

29 assert StepType.SEQUENTIAL.value == "sequential" 

30 assert StepType.PARALLEL.value == "parallel" 

31 assert StepType.CONDITIONAL.value == "conditional" 

32 assert StepType.LOOP.value == "loop" 

33 assert StepType.SUB_WORKFLOW.value == "sub" 

34 assert StepType.JOIN.value == "join" 

35 assert StepType.SPLIT.value == "split" 

36 

37 

38class TestExecutionStatus: 

39 def test_values(self): 

40 assert ExecutionStatus.PENDING.value == "pending" 

41 assert ExecutionStatus.RUNNING.value == "running" 

42 assert ExecutionStatus.SUCCESS.value == "success" 

43 assert ExecutionStatus.FAILED.value == "failed" 

44 assert ExecutionStatus.SKIPPED.value == "skipped" 

45 assert ExecutionStatus.CANCELLED.value == "cancelled" 

46 assert ExecutionStatus.RETRYING.value == "retrying" 

47 

48 

49class TestErrorStrategy: 

50 def test_values(self): 

51 assert ErrorStrategy.RETRY.value == "retry" 

52 assert ErrorStrategy.FALLBACK.value == "fallback" 

53 assert ErrorStrategy.SKIP.value == "skip" 

54 assert ErrorStrategy.ESCALATE.value == "escalate" 

55 assert ErrorStrategy.PAUSE.value == "pause" 

56 

57 

58class TestConditionOperator: 

59 def test_values(self): 

60 assert ConditionOperator.EQUALS.value == "eq" 

61 assert ConditionOperator.NOT_EQUALS.value == "neq" 

62 assert ConditionOperator.CONTAINS.value == "contains" 

63 assert ConditionOperator.GREATER.value == "gt" 

64 assert ConditionOperator.LESS.value == "lt" 

65 assert ConditionOperator.IN.value == "in" 

66 assert ConditionOperator.MATCHES.value == "matches" 

67 assert ConditionOperator.EXISTS.value == "exists" 

68 assert ConditionOperator.EMPTY.value == "empty" 

69 

70 

71# ============================================================================ 

72# WorkflowContext 

73# ============================================================================ 

74 

75class TestWorkflowContext: 

76 def test_default_construction(self): 

77 ctx = WorkflowContext() 

78 assert ctx.variables == {} 

79 assert ctx.history == [] 

80 assert ctx.errors == [] 

81 

82 def test_get_simple_key(self): 

83 ctx = WorkflowContext(variables={"a": 1, "b": "hello"}) 

84 assert ctx.get("a") == 1 

85 assert ctx.get("b") == "hello" 

86 

87 def test_get_missing_with_default(self): 

88 ctx = WorkflowContext() 

89 assert ctx.get("missing", "default") == "default" 

90 assert ctx.get("missing") is None 

91 

92 def test_get_dot_notation(self): 

93 ctx = WorkflowContext(variables={"result": {"output": {"text": "hello"}}}) 

94 assert ctx.get("result.output.text") == "hello" 

95 

96 def test_get_dot_notation_missing_mid(self): 

97 ctx = WorkflowContext(variables={"a": 1}) 

98 assert ctx.get("a.b.c", "fallback") == "fallback" 

99 

100 def test_get_non_dict_intermediary(self): 

101 ctx = WorkflowContext(variables={"a": 1}) 

102 assert ctx.get("a.b", "fallback") == "fallback" 

103 

104 def test_set_simple_key(self): 

105 ctx = WorkflowContext() 

106 ctx.set("x", 42) 

107 assert ctx.variables == {"x": 42} 

108 

109 def test_set_dot_notation_creates_nested(self): 

110 ctx = WorkflowContext() 

111 ctx.set("a.b.c", 99) 

112 assert ctx.variables == {"a": {"b": {"c": 99}}} 

113 

114 def test_set_dot_notation_extends_existing(self): 

115 ctx = WorkflowContext(variables={"a": {"b": 1}}) 

116 ctx.set("a.c", 2) 

117 assert ctx.variables == {"a": {"b": 1, "c": 2}} 

118 

119 

120# ============================================================================ 

121# StepResult 

122# ============================================================================ 

123 

124class TestStepResult: 

125 def test_defaults(self): 

126 r = StepResult("step1", ExecutionStatus.SUCCESS) 

127 assert r.step_id == "step1" 

128 assert r.status == ExecutionStatus.SUCCESS 

129 assert r.output is None 

130 assert r.error is None 

131 assert r.duration == 0.0 

132 assert r.retries == 0 

133 

134 def test_full_fields(self): 

135 r = StepResult( 

136 "s1", ExecutionStatus.FAILED, output="partial", 

137 error="boom", duration=1.5, retries=2, metadata={"k": "v"}, 

138 ) 

139 assert r.output == "partial" 

140 assert r.error == "boom" 

141 assert r.duration == 1.5 

142 assert r.retries == 2 

143 assert r.metadata == {"k": "v"} 

144 

145 

146# ============================================================================ 

147# WorkflowStep 

148# ============================================================================ 

149 

150class TestWorkflowStep: 

151 def test_minimal(self): 

152 s = WorkflowStep(id="s1", type=StepType.TASK) 

153 assert s.id == "s1" 

154 assert s.type == StepType.TASK 

155 assert s.name == "" 

156 assert s.children == [] 

157 

158 def test_with_agent_and_task(self): 

159 s = WorkflowStep( 

160 id="greet", type=StepType.TASK, agent="greeter", task="Say hello" 

161 ) 

162 assert s.agent == "greeter" 

163 assert s.task == "Say hello" 

164 

165 def test_default_values(self): 

166 s = WorkflowStep(id="s1", type=StepType.TASK) 

167 assert s.max_retries == 3 

168 assert s.retry_delay == 1.0 

169 assert s.timeout == 300.0 

170 assert s.on_error == ErrorStrategy.ESCALATE 

171 assert s.depends_on == [] 

172 

173 

174# ============================================================================ 

175# WorkflowDefinition 

176# ============================================================================ 

177 

178class TestWorkflowDefinition: 

179 def test_empty_construction(self): 

180 wf = WorkflowDefinition(name="test") 

181 assert wf.name == "test" 

182 assert wf.version == "1.0" 

183 assert wf.root is None 

184 

185 def test_steps_property_empty(self): 

186 wf = WorkflowDefinition(name="empty") 

187 assert wf.steps == [] 

188 

189 def test_steps_property_flat_collection(self): 

190 s1 = WorkflowStep(id="a", type=StepType.TASK) 

191 s2 = WorkflowStep(id="b", type=StepType.TASK) 

192 s1.children = [s2] 

193 wf = WorkflowDefinition(name="test", root=s1) 

194 step_ids = [s.id for s in wf.steps] 

195 assert step_ids == ["a", "b"] 

196 

197 def test_steps_property_deep_nesting(self): 

198 s3 = WorkflowStep(id="c", type=StepType.TASK) 

199 s2 = WorkflowStep(id="b", type=StepType.SEQUENTIAL, children=[s3]) 

200 s1 = WorkflowStep(id="a", type=StepType.SEQUENTIAL, children=[s2]) 

201 wf = WorkflowDefinition(name="deep", root=s1) 

202 assert [s.id for s in wf.steps] == ["a", "b", "c"] 

203 

204 def test_validate_no_root(self): 

205 wf = WorkflowDefinition(name="bad") 

206 issues = wf.validate() 

207 assert "no root step" in issues[0] 

208 

209 def test_validate_duplicate_ids(self): 

210 s1 = WorkflowStep(id="dup", type=StepType.TASK) 

211 s2 = WorkflowStep(id="dup", type=StepType.TASK) 

212 s1.children = [s2] 

213 wf = WorkflowDefinition(name="test", root=s1) 

214 issues = wf.validate() 

215 assert any("Duplicate step ID" in i for i in issues) 

216 

217 def test_validate_conditional_no_condition(self): 

218 s = WorkflowStep(id="cond", type=StepType.CONDITIONAL) 

219 wf = WorkflowDefinition(name="test", root=s) 

220 issues = wf.validate() 

221 assert any("no condition" in i for i in issues) 

222 

223 def test_validate_task_no_agent(self): 

224 s = WorkflowStep(id="t1", type=StepType.TASK) 

225 wf = WorkflowDefinition(name="test", root=s) 

226 issues = wf.validate() 

227 assert any("no agent assigned" in i for i in issues) 

228 

229 def test_validate_unknown_depends_on(self): 

230 s = WorkflowStep(id="t1", type=StepType.TASK, agent="a", depends_on=["unknown"]) 

231 wf = WorkflowDefinition(name="test", root=s) 

232 issues = wf.validate() 

233 assert any("depends on unknown step" in i for i in issues) 

234 

235 def test_validate_clean(self): 

236 s = WorkflowStep(id="t1", type=StepType.TASK, agent="a") 

237 wf = WorkflowDefinition(name="test", root=s) 

238 issues = wf.validate() 

239 assert issues == [] 

240 

241 def test_to_mermaid_task(self): 

242 s = WorkflowStep(id="greet", type=StepType.TASK, name="Greeter") 

243 wf = WorkflowDefinition(name="mermaid_test", root=s) 

244 result = wf.to_mermaid() 

245 assert "graph TD" in result 

246 assert "greet" in result 

247 assert "Greeter" in result 

248 

249 def test_to_mermaid_sequential(self): 

250 s2 = WorkflowStep(id="b", type=StepType.TASK, name="Step B") 

251 s1 = WorkflowStep(id="a", type=StepType.SEQUENTIAL, name="Step A", children=[s2]) 

252 wf = WorkflowDefinition(name="mermaid_seq", root=s1) 

253 result = wf.to_mermaid() 

254 assert "a --> b" in result 

255 

256 def test_to_mermaid_conditional(self): 

257 true_s = WorkflowStep(id="yes", type=StepType.TASK) 

258 false_s = WorkflowStep(id="no", type=StepType.TASK) 

259 s = WorkflowStep( 

260 id="check", type=StepType.CONDITIONAL, 

261 branches={"true": [true_s], "false": [false_s]}, 

262 ) 

263 wf = WorkflowDefinition(name="mermaid_cond", root=s) 

264 result = wf.to_mermaid() 

265 assert "yes" in result 

266 assert "no" in result 

267 

268 def test_to_mermaid_no_root(self): 

269 wf = WorkflowDefinition(name="empty") 

270 result = wf.to_mermaid() 

271 assert result == "graph TD" 

272 

273 

274# ============================================================================ 

275# ConditionEvaluator 

276# ============================================================================ 

277 

278class TestConditionEvaluator: 

279 def test_empty_condition(self): 

280 assert ConditionEvaluator.evaluate({}, WorkflowContext()) is True 

281 assert ConditionEvaluator.evaluate(None, WorkflowContext()) is True 

282 

283 def test_equals(self): 

284 ctx = WorkflowContext(variables={"x": 5}) 

285 assert ConditionEvaluator.evaluate({"field": "x", "op": "eq", "value": 5}, ctx) is True 

286 assert ConditionEvaluator.evaluate({"field": "x", "op": "eq", "value": 3}, ctx) is False 

287 

288 def test_not_equals(self): 

289 ctx = WorkflowContext(variables={"x": 5}) 

290 assert ConditionEvaluator.evaluate({"field": "x", "op": "neq", "value": 3}, ctx) is True 

291 assert ConditionEvaluator.evaluate({"field": "x", "op": "neq", "value": 5}, ctx) is False 

292 

293 def test_contains(self): 

294 ctx = WorkflowContext(variables={"text": "hello world"}) 

295 assert ConditionEvaluator.evaluate( 

296 {"field": "text", "op": "contains", "value": "world"}, ctx 

297 ) is True 

298 assert ConditionEvaluator.evaluate( 

299 {"field": "text", "op": "contains", "value": "xyz"}, ctx 

300 ) is False 

301 

302 def test_greater(self): 

303 ctx = WorkflowContext(variables={"x": 10}) 

304 assert ConditionEvaluator.evaluate({"field": "x", "op": "gt", "value": 5}, ctx) is True 

305 assert ConditionEvaluator.evaluate({"field": "x", "op": "gt", "value": 15}, ctx) is False 

306 

307 def test_less(self): 

308 ctx = WorkflowContext(variables={"x": 5}) 

309 assert ConditionEvaluator.evaluate({"field": "x", "op": "lt", "value": 10}, ctx) is True 

310 assert ConditionEvaluator.evaluate({"field": "x", "op": "lt", "value": 3}, ctx) is False 

311 

312 def test_in(self): 

313 ctx = WorkflowContext(variables={"x": "a"}) 

314 assert ConditionEvaluator.evaluate( 

315 {"field": "x", "op": "in", "value": ["a", "b", "c"]}, ctx 

316 ) is True 

317 assert ConditionEvaluator.evaluate( 

318 {"field": "x", "op": "in", "value": ["d", "e"]}, ctx 

319 ) is False 

320 

321 def test_matches(self): 

322 ctx = WorkflowContext(variables={"email": "user@test.com"}) 

323 assert ConditionEvaluator.evaluate( 

324 {"field": "email", "op": "matches", "value": r".*@test\.com"}, ctx 

325 ) is True 

326 assert ConditionEvaluator.evaluate( 

327 {"field": "email", "op": "matches", "value": r".*@other\.com"}, ctx 

328 ) is False 

329 

330 def test_exists(self): 

331 ctx = WorkflowContext(variables={"x": 1}) 

332 assert ConditionEvaluator.evaluate({"field": "x", "op": "exists"}, ctx) is True 

333 assert ConditionEvaluator.evaluate({"field": "missing", "op": "exists"}, ctx) is False 

334 

335 def test_empty(self): 

336 ctx = WorkflowContext(variables={"x": "", "y": "hello", "z": None}) 

337 assert ConditionEvaluator.evaluate({"field": "x", "op": "empty"}, ctx) is True 

338 assert ConditionEvaluator.evaluate({"field": "y", "op": "empty"}, ctx) is False 

339 assert ConditionEvaluator.evaluate({"field": "z", "op": "empty"}, ctx) is True 

340 

341 def test_and_combinator(self): 

342 ctx = WorkflowContext(variables={"a": 1, "b": 2}) 

343 cond = {"and": [ 

344 {"field": "a", "op": "eq", "value": 1}, 

345 {"field": "b", "op": "gt", "value": 0}, 

346 ]} 

347 assert ConditionEvaluator.evaluate(cond, ctx) is True 

348 cond2 = {"and": [ 

349 {"field": "a", "op": "eq", "value": 1}, 

350 {"field": "b", "op": "gt", "value": 10}, 

351 ]} 

352 assert ConditionEvaluator.evaluate(cond2, ctx) is False 

353 

354 def test_or_combinator(self): 

355 ctx = WorkflowContext(variables={"a": 1, "b": 2}) 

356 cond = {"or": [ 

357 {"field": "a", "op": "eq", "value": 99}, 

358 {"field": "b", "op": "eq", "value": 2}, 

359 ]} 

360 assert ConditionEvaluator.evaluate(cond, ctx) is True 

361 cond2 = {"or": [ 

362 {"field": "a", "op": "eq", "value": 99}, 

363 {"field": "b", "op": "eq", "value": 99}, 

364 ]} 

365 assert ConditionEvaluator.evaluate(cond2, ctx) is False 

366 

367 def test_not_combinator(self): 

368 ctx = WorkflowContext(variables={"x": 5}) 

369 cond = {"not": {"field": "x", "op": "eq", "value": 3}} 

370 assert ConditionEvaluator.evaluate(cond, ctx) is True 

371 

372 def test_contains_none_field(self): 

373 ctx = WorkflowContext() 

374 assert ConditionEvaluator.evaluate( 

375 {"field": "missing", "op": "contains", "value": "x"}, ctx 

376 ) is False 

377 

378 def test_gt_lt_invalid_types(self): 

379 ctx = WorkflowContext(variables={"x": "abc"}) 

380 assert ConditionEvaluator.evaluate( 

381 {"field": "x", "op": "gt", "value": 5}, ctx 

382 ) is False 

383 assert ConditionEvaluator.evaluate( 

384 {"field": "x", "op": "lt", "value": 5}, ctx 

385 ) is False 

386 

387 def test_in_non_iterable_value(self): 

388 ctx = WorkflowContext(variables={"x": "a"}) 

389 assert ConditionEvaluator.evaluate( 

390 {"field": "x", "op": "in", "value": "not_a_list"}, ctx 

391 ) is False 

392 

393 def test_matches_invalid_regex(self): 

394 ctx = WorkflowContext(variables={"x": "test"}) 

395 assert ConditionEvaluator.evaluate( 

396 {"field": "x", "op": "matches", "value": "["}, ctx 

397 ) is False 

398 

399 

400# ============================================================================ 

401# WorkflowParser 

402# ============================================================================ 

403 

404class TestWorkflowParser: 

405 def test_parse_str_json(self): 

406 text = json.dumps({"name": "test_wf", "steps": [ 

407 {"id": "s1", "type": "task", "agent": "gpt"} 

408 ]}) 

409 wf = WorkflowParser.parse_str(text) 

410 assert wf.name == "test_wf" 

411 assert wf.root.id == "s1" 

412 

413 def test_parse_str_yaml(self): 

414 text = """name: test_wf 

415steps: 

416 - id: s1 

417 type: task 

418 agent: gpt 

419""" 

420 wf = WorkflowParser.parse_str(text) 

421 assert wf.name == "test_wf" 

422 assert wf.root.id == "s1" 

423 

424 def test_parse_dict_basic(self): 

425 wf = WorkflowParser.parse_dict({ 

426 "name": "basic", 

427 "steps": [{"id": "s1", "type": "task", "agent": "a"}], 

428 }) 

429 assert wf.name == "basic" 

430 assert wf.root.id == "s1" 

431 assert wf.root.agent == "a" 

432 

433 def test_parse_dict_multiple_steps_chained(self): 

434 wf = WorkflowParser.parse_dict({ 

435 "name": "chain", 

436 "steps": [ 

437 {"id": "s1", "type": "task", "agent": "a"}, 

438 {"id": "s2", "type": "task", "agent": "b"}, 

439 {"id": "s3", "type": "task", "agent": "c"}, 

440 ], 

441 }) 

442 # s1.children = [s2], s2.children = [s3] 

443 assert wf.root.id == "s1" 

444 assert len(wf.root.children) == 1 

445 assert wf.root.children[0].id == "s2" 

446 assert wf.root.children[0].children[0].id == "s3" 

447 

448 def test_parse_dict_no_steps_raises(self): 

449 with pytest.raises(ValueError, match="No steps defined"): 

450 WorkflowParser.parse_dict({"name": "bad", "steps": []}) 

451 

452 def test_parse_dict_with_branches(self): 

453 wf = WorkflowParser.parse_dict({ 

454 "name": "branchy", 

455 "steps": [{ 

456 "id": "cond", 

457 "type": "conditional", 

458 "condition": {"field": "x", "op": "eq", "value": 1}, 

459 "branches": { 

460 "true": [{"id": "yes", "type": "task", "agent": "a"}], 

461 "false": [{"id": "no", "type": "task", "agent": "b"}], 

462 }, 

463 }], 

464 }) 

465 assert wf.root.id == "cond" 

466 assert len(wf.root.branches["true"]) == 1 

467 assert wf.root.branches["true"][0].id == "yes" 

468 

469 def test_parse_dict_with_children(self): 

470 wf = WorkflowParser.parse_dict({ 

471 "name": "nested", 

472 "steps": [{ 

473 "id": "root", 

474 "type": "sequential", 

475 "children": [ 

476 {"id": "a", "type": "task", "agent": "x"}, 

477 {"id": "b", "type": "task", "agent": "y"}, 

478 ], 

479 }], 

480 }) 

481 assert wf.root.id == "root" 

482 assert len(wf.root.children) == 2 

483 

484 def test_parse_dict_loop(self): 

485 wf = WorkflowParser.parse_dict({ 

486 "name": "loopy", 

487 "steps": [{ 

488 "id": "loop", 

489 "type": "loop", 

490 "max_iterations": 5, 

491 "loop_condition": {"field": "done", "op": "eq", "value": True}, 

492 "children": [{"id": "iter", "type": "task", "agent": "a"}], 

493 }], 

494 }) 

495 assert wf.root.type == StepType.LOOP 

496 assert wf.root.max_iterations == 5 

497 

498 def test_parse_dict_error_strategies(self): 

499 wf = WorkflowParser.parse_dict({ 

500 "name": "errors", 

501 "steps": [{ 

502 "id": "s1", "type": "task", "agent": "a", 

503 "on_error": "retry", "max_retries": 5, "retry_delay": 3.0, 

504 "fallback": {"id": "fb", "type": "task", "agent": "fb_agent"}, 

505 }], 

506 }) 

507 assert wf.root.on_error == ErrorStrategy.RETRY 

508 assert wf.root.max_retries == 5 

509 assert wf.root.retry_delay == 3.0 

510 assert wf.root.fallback_step is not None 

511 assert wf.root.fallback_step.id == "fb" 

512 

513 def test_to_yaml(self): 

514 wf = WorkflowDefinition(name="test", version="1.0") 

515 wf.root = WorkflowStep(id="s1", type=StepType.TASK, agent="a") 

516 yaml_str = WorkflowParser.to_yaml(wf) 

517 assert "name: test" in yaml_str 

518 assert "id: s1" in yaml_str 

519 

520 def test_to_json(self): 

521 wf = WorkflowDefinition(name="test") 

522 wf.root = WorkflowStep(id="s1", type=StepType.TASK, agent="a") 

523 json_str = WorkflowParser.to_json(wf) 

524 data = json.loads(json_str) 

525 assert data["name"] == "test" 

526 assert data["steps"][0]["id"] == "s1" 

527 

528 def test_roundtrip_yaml(self): 

529 original = WorkflowParser.parse_dict({ 

530 "name": "roundtrip", 

531 "steps": [{"id": "s1", "type": "task", "agent": "gpt"}], 

532 }) 

533 yaml_str = WorkflowParser.to_yaml(original) 

534 parsed = WorkflowParser.parse_str(yaml_str) 

535 assert parsed.name == original.name 

536 assert parsed.root.id == original.root.id 

537 

538 def test_roundtrip_json(self): 

539 original = WorkflowParser.parse_dict({ 

540 "name": "roundtrip", 

541 "steps": [{"id": "s1", "type": "task", "agent": "gpt"}], 

542 }) 

543 json_str = WorkflowParser.to_json(original) 

544 parsed = WorkflowParser.parse_str(json_str) 

545 assert parsed.name == original.name 

546 assert parsed.root.id == original.root.id 

547 

548 def test_parse_file_yaml(self, tmp_path): 

549 yaml_path = tmp_path / "test.yaml" 

550 yaml_path.write_text("""name: file_test 

551steps: 

552 - id: s1 

553 type: task 

554 agent: gpt 

555""") 

556 wf = WorkflowParser.parse_file(str(yaml_path)) 

557 assert wf.name == "file_test" 

558 assert wf.root.id == "s1" 

559 

560 def test_parse_file_json(self, tmp_path): 

561 json_path = tmp_path / "test.json" 

562 json_path.write_text(json.dumps({ 

563 "name": "file_test", 

564 "steps": [{"id": "s1", "type": "task", "agent": "gpt"}], 

565 })) 

566 wf = WorkflowParser.parse_file(str(json_path)) 

567 assert wf.name == "file_test" 

568 

569 

570# ============================================================================ 

571# WorkflowTemplates 

572# ============================================================================ 

573 

574class TestWorkflowTemplates: 

575 def test_sequential(self): 

576 wf = WorkflowTemplates.sequential("seq", ["a", "b", "c"], "Task: {{input}}") 

577 assert wf.name == "seq" 

578 assert wf.root.id == "step_a" 

579 assert wf.root.children[0].id == "step_b" 

580 assert wf.root.children[0].children[0].id == "step_c" 

581 

582 def test_parallel_broadcast(self): 

583 wf = WorkflowTemplates.parallel_broadcast("broad", ["a", "b"], "process") 

584 assert wf.root.type == StepType.PARALLEL 

585 assert len(wf.root.children) == 2 

586 

587 def test_map_reduce(self): 

588 wf = WorkflowTemplates.map_reduce( 

589 "mr", ["m1", "m2"], "r1", "map task", "reduce task" 

590 ) 

591 assert wf.root.type == StepType.PARALLEL 

592 assert wf.root.id == "map_phase" 

593 assert len(wf.root.children) == 1 

594 assert wf.root.children[0].id == "reduce_phase" 

595 

596 def test_conditional_branch(self): 

597 wf = WorkflowTemplates.conditional_branch( 

598 "cond", "score", "good", "bad", "do: {{score}}" 

599 ) 

600 assert wf.root.type == StepType.CONDITIONAL 

601 assert "true" in wf.root.branches 

602 assert "false" in wf.root.branches 

603 

604 def test_retry_loop(self): 

605 wf = WorkflowTemplates.retry_loop("retry", "agent1", "task", max_retries=5) 

606 assert wf.root.on_error == ErrorStrategy.RETRY 

607 assert wf.root.max_retries == 5 

608 assert wf.root.retry_delay == 2.0 

609 

610 

611# ============================================================================ 

612# WorkflowEngine 

613# ============================================================================ 

614 

615class TestWorkflowEngine: 

616 @pytest.mark.asyncio 

617 async def test_execute_single_task(self): 

618 wf = WorkflowDefinition( 

619 name="simple", 

620 root=WorkflowStep(id="t1", type=StepType.TASK, agent="gpt", task="hello"), 

621 ) 

622 engine = WorkflowEngine() 

623 ctx = await engine.execute(wf) 

624 assert ctx.variables["steps"]["t1"]["output"] is not None 

625 

626 @pytest.mark.asyncio 

627 async def test_execute_sequential(self): 

628 s2 = WorkflowStep(id="t2", type=StepType.TASK, agent="b", task="step2") 

629 s1 = WorkflowStep(id="t1", type=StepType.SEQUENTIAL, children=[s2]) 

630 wf = WorkflowDefinition(name="seq", root=s1) 

631 engine = WorkflowEngine() 

632 ctx = await engine.execute(wf) 

633 assert ctx.variables["steps"]["t2"]["output"] is not None 

634 

635 @pytest.mark.asyncio 

636 async def test_execute_parallel(self): 

637 children = [ 

638 WorkflowStep(id="a", type=StepType.TASK, agent="a", task="a"), 

639 WorkflowStep(id="b", type=StepType.TASK, agent="b", task="b"), 

640 ] 

641 root = WorkflowStep(id="par", type=StepType.PARALLEL, children=children) 

642 wf = WorkflowDefinition(name="par", root=root) 

643 engine = WorkflowEngine() 

644 ctx = await engine.execute(wf) 

645 outputs = ctx.variables["steps"]["par"]["outputs"] 

646 assert "a" in outputs 

647 assert "b" in outputs 

648 

649 @pytest.mark.asyncio 

650 async def test_execute_conditional_true(self): 

651 true_s = WorkflowStep(id="yes", type=StepType.TASK, agent="a", task="true path") 

652 false_s = WorkflowStep(id="no", type=StepType.TASK, agent="b", task="false path") 

653 root = WorkflowStep( 

654 id="check", 

655 type=StepType.CONDITIONAL, 

656 condition={"field": "flag", "op": "eq", "value": True}, 

657 branches={"true": [true_s], "false": [false_s]}, 

658 ) 

659 wf = WorkflowDefinition(name="cond", root=root, variables={"flag": True}) 

660 engine = WorkflowEngine() 

661 ctx = await engine.execute(wf) 

662 # Check that the true branch was executed 

663 assert ctx.variables["steps"]["yes"]["output"] is not None 

664 # Confirm false branch was NOT executed 

665 assert ctx.get("steps.no.output") is None 

666 

667 @pytest.mark.asyncio 

668 async def test_execute_conditional_false(self): 

669 true_s = WorkflowStep(id="yes", type=StepType.TASK, agent="a", task="true path")