def test_py_to_daghandler__sanity(self):
from .data.valid import valid_dag
# Double passthrough - as default args are populated when
# instantiating dag
handler = DagHandler.load_from_dag(valid_dag)
handler_pycode = handler.to_python()
# Recreate handler
mod = import_str_as_module(handler_pycode, "test_py2dag")
handler = DagHandler.load_from_dag(getattr(mod, "valid_dag"))
assert handler.to_python() == handler_pycode
def test_wml_marshall__fails_on_missing_required_param(self):
new_wml = deepcopy(self.valid_wml)
ind = [
i for i, val in enumerate(new_wml["dag"]["parameters"])
if val["id"] == "dag_id"
]
assert len(ind) == 1
ind = ind[0]
new_wml["dag"]["parameters"][ind].pop("value")
with self.assertRaises(DagHandlerValidationError):
DagHandler.load_from_wml(new_wml)
def test_wml_marshall__fails_on_circular_tasks(self):
# Add in an additional link that closes the loop - this should break the test
existing_link = deepcopy(self.valid_wml_dict)["links"].popitem()[1]
self.valid_wml_dict["links"]["a cyclical link"] = {
"id": "a cyclical link",
"from_node": {
"nodeId": existing_link["to_node"]["nodeId"]
},
"to_node": {
"nodeId": existing_link["from_node"]["nodeId"]
},
}
with self.assertRaises(DagHandlerValidationError):
DagHandler.load_from_wml(self.valid_wml_dict)
def test_wml_conversion_to_python__fails_on_no_start_date(self):
new_wml = deepcopy(self.valid_wml_dict)
for param in new_wml["dag"]["parameters"]:
if param["id"] == "start_date":
param.pop("value")
dag_handler = DagHandler.load_from_wml(new_wml)
with self.assertRaises(DagHandlerValidationError):
dag_handler.to_python()
def test_dag_params_from_wml(self):
di = DagHandler.load_from_wml(self.valid_wml_dict)
assert di.params["dag_id"]["value"] == "ValidDag"
assert [t.operator_type
for t in di.tasks] == ["BashOperator", "BashOperator"]
assert len(di.links.node_ids) == 2
# Note that the rest of the values are defaults and are ignored
assert list(di.params.keys()) == ["dag_id", "description"]
def test_wml_conversion_to_python__succeeds_on_start_date_only_in_tasks(
self):
new_wml = deepcopy(self.valid_wml_dict)
for param in new_wml["dag"]["parameters"]:
if param["id"] == "start_date":
param.pop("value")
for node in new_wml["nodes"].values():
for param in node["properties"]["parameters"]:
if param["id"] == "start_date":
param["value"] = "2020/05/20"
dag_handler = DagHandler.load_from_wml(new_wml)
py_code = dag_handler.to_python()
assert py_code
def test_wml_conversion_to_python__fails_on_start_date_missing_from_one_task(
self):
new_wml = deepcopy(self.valid_wml_dict)
for param in new_wml["dag"]["parameters"]:
if param["id"] == "start_date":
param.pop("value")
node_ids = list(new_wml["nodes"].keys())
for param in new_wml["nodes"][node_ids[0]]["properties"]["parameters"]:
if param["id"] == "start_date":
param["value"] = "2020/05/20"
dag_handler = DagHandler.load_from_wml(new_wml)
with self.assertRaises(DagHandlerValidationError):
dag_handler.to_python()
def test_dag_handler_to_wml__sanity(self):
in_d = MinimalWmlSchema().dump(self.valid_wml_dict)
dag_handler = DagHandler.load_from_wml(self.valid_wml_dict)
out_d = dag_handler.to_wml()
# Only differences should be
# - Link IDs aren't retained across conversion
# - Values that match defaults are skipped from Py -> WML
# - Position values will differ
# Tidy up links Dict
assert len(in_d["links"]) == len(out_d["links"])
assert len(in_d["links"]) == 1
links_dict = out_d["links"].pop(list(out_d["links"].keys())[0])
in_d_links_id = list(in_d["links"].keys())[0]
links_dict["id"] = in_d_links_id
out_d["links"][in_d_links_id] = links_dict
self.assertDictEqual(in_d["links"], out_d["links"])
# Remove any values from in dict where value == default
for param in in_d["dag"]["parameters"]:
if param.get("value", "") == param.get("default", None):
param.pop("value")
self.assertDictEqual(in_d["dag"], out_d["dag"])
for node in in_d["nodes"].values():
node["position"]["x"] = 0
node["position"]["y"] = 0
for param in node["properties"]["parameters"]:
if param.get("value", "") == param.get("default", None):
param.pop("value")
for node in out_d["nodes"].values():
node["position"]["x"] = 0
node["position"]["y"] = 0
self.assertDictEqual(in_d["nodes"], out_d["nodes"])
def test_wml_conversion_to_python__sanity(self):
dag_handler = DagHandler.load_from_wml(self.valid_wml_dict)
res = dag_handler.to_python()
print(res)
assert res.strip() == self.valid_py.strip()