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()