|
| 1 | +import unittest |
| 2 | +import os |
| 3 | +from pyiron_workflow import Workflow, to_function_node |
| 4 | +from python_workflow_definition.pyiron_workflow import load_workflow_json, write_workflow_json |
| 5 | + |
| 6 | +function_str = """ |
| 7 | +def get_prod_and_div(x, y): |
| 8 | + return {"prod": x * y, "div": x / y} |
| 9 | +
|
| 10 | +
|
| 11 | +def get_sum(x, y): |
| 12 | + return x + y |
| 13 | +
|
| 14 | +
|
| 15 | +def get_square(x): |
| 16 | + return x ** 2 |
| 17 | +""" |
| 18 | + |
| 19 | + |
| 20 | +class TestPyironWorkflow(unittest.TestCase): |
| 21 | + def test_pyiron_workflow(self): |
| 22 | + workflow_json_filename = "pyiron_workflow_arithmetic.json" |
| 23 | + with open("workflow.py", "w") as f: |
| 24 | + f.write(function_str) |
| 25 | + |
| 26 | + from workflow import get_prod_and_div as _get_prod_and_div |
| 27 | + from workflow import get_sum as _get_sum |
| 28 | + from workflow import get_square as _get_square |
| 29 | + |
| 30 | + get_prod_and_div = to_function_node("get_prod_and_div", _get_prod_and_div, "get_prod_and_div") |
| 31 | + get_sum = to_function_node("get_sum", _get_sum, "get_sum") |
| 32 | + get_square = to_function_node("get_square", _get_square, "get_square") |
| 33 | + wf = Workflow("my_workflow") |
| 34 | + wf.x = 1 |
| 35 | + wf.y = 2 |
| 36 | + wf.prod_and_div = get_prod_and_div(x=wf.x, y=wf.y) |
| 37 | + wf.tmp_sum = get_sum(x=wf.prod_and_div["prod"], y=wf.prod_and_div["div"]) |
| 38 | + wf.square_result = get_square(x=wf.tmp_sum) |
| 39 | + write_workflow_json(graph_as_dict=wf.graph_as_dict, file_name=workflow_json_filename) |
| 40 | + wf = load_workflow_json(file_name=workflow_json_filename) |
| 41 | + wf.run() |
| 42 | + |
| 43 | + self.assertTrue(os.path.exists(workflow_json_filename)) |
0 commit comments