def test_grupnet_rst_docs(tmp_path): """Provide the input and output for the examples in the RST documentation""" os.chdir(tmp_path) schstr = """ START 01 'JAN' 2000 / SCHEDULE GRUPTREE 'OPEAST' 'OP' / 'OPWEST' 'OP' / 'INJEAST' 'WI' / 'OP' 'FIELD' / 'WI' 'FIELD' / 'FIELD' 'AREA' / 'AREA' 'NORTHSEA' / / GRUPNET 'FIELD' 90 / 'OPWEST' 100 / / WELSPECS 'OP1' 'OPWEST' 41 125 1759.74 'OIL' 0.0 'STD' 'SHUT' 'YES' 0 'SEG' / 'OP2' 'OPEAST' 43 122 1776.01 'OIL' 0.0 'STD' 'SHUT' 'YES' 0 'SEG' / 'INJ1' 'INJEAST' 33 115 1960.21 'OIL' 0.0 'STD' 'SHUT' 'YES' 0 'SEG' / / """ deck = EclFiles.str2deck(schstr) grupdf = gruptree.df(deck) grupdf[["DATE", "CHILD", "PARENT", "KEYWORD"]].to_csv("gruptree.csv", index=False) grupdf.to_csv("gruptreenet.csv", index=False) grup_dict = gruptree.edge_dataframe2dict(grupdf) print("Copy and paste into RST files:") print(str(gruptree.tree_from_dict(grup_dict[0]))) assert ( str(gruptree.tree_from_dict(grup_dict[0])).strip() == """ NORTHSEA └── AREA └── FIELD ├── OP │ ├── OPEAST │ │ └── OP2 │ └── OPWEST │ └── OP1 └── WI └── INJEAST └── INJ1 """.strip() )
def test_emptytree(): """Test empty schedule sections. Don't want to crash""" schstr = "" deck = EclFiles.str2deck(schstr) grupdf = gruptree.df(deck) assert grupdf.empty gruptreedict = gruptree.edge_dataframe2dict(grupdf) assert not gruptreedict[0] treelibtree = gruptree.dict2treelib("", gruptreedict[0]) treestring = str(treelibtree) assert not treestring.strip() # Let it return whitespace
def test_emptytree_strdeck(): """Test empty schedule sections. Don't want to crash""" schstr = "" deck = EclFiles.str2deck(schstr) grupdf = gruptree.df(deck) assert grupdf.empty gruptreedict = gruptree.edge_dataframe2dict(grupdf) assert not gruptreedict[0] treelibtree = gruptree.tree_from_dict(gruptreedict[0]) # Returning an empty string and not a treelib.Tree() is # a workaround for a limitation in treelib. assert treelibtree == ""
def test_multiple_roots(): """Test edge_dataframe2dict with multiple roots""" answer = [ {"FIELDA": {"PLATA": {}}}, {"FIELDB": {"PLATB": {}}}, ] edges = pd.DataFrame( [ {"CHILD": "FIELDA", "PARENT": None}, {"CHILD": "FIELDB", "PARENT": None}, {"CHILD": "PLATA", "PARENT": "FIELDA"}, {"CHILD": "PLATB", "PARENT": "FIELDB"}, ] ) assert gruptree.edge_dataframe2dict(edges) == answer # Same result if the dummy rows for the roots are omitted: edges_noroots = pd.DataFrame( [ {"CHILD": "PLATA", "PARENT": "FIELDA"}, {"CHILD": "PLATB", "PARENT": "FIELDB"}, ] ) assert gruptree.edge_dataframe2dict(edges_noroots) == answer # And order does not matter, should be sorted on root node label: edges_noroots = pd.DataFrame( [ {"CHILD": "PLATB", "PARENT": "FIELDB"}, {"CHILD": "PLATA", "PARENT": "FIELDA"}, ] ) assert gruptree.edge_dataframe2dict(edges_noroots) == answer # The function tree_from_dict should be called with one tree at a time: with pytest.raises(ValueError, match="single tree"): gruptree.tree_from_dict({"foo": 1, "bar": 2})
def test_grupnetroot(schstr, expected_dframe, expected_tree): """Test that terminal pressure of the tree root can be included in the dataframe (with an empty parent)""" deck = EclFiles.str2deck(schstr) grupdf = gruptree.df(deck, startdate="2000-01-01") non_default_columns = ["CHILD", "PARENT", "TERMINAL_PRESSURE"] pd.testing.assert_frame_equal( grupdf[non_default_columns] .sort_values(["CHILD", "PARENT"]) .reset_index(drop=True), expected_dframe.sort_values(["CHILD", "PARENT"]).reset_index(drop=True), check_dtype=False, ) treelist = gruptree.edge_dataframe2dict(grupdf) # Merge strings for all trees (if multiple roots) strtrees = [str(gruptree.tree_from_dict(tree)) for tree in treelist] strtrees.sort() # Avoid flaky test due to sorting treelibtree = "".join(strtrees) assert treelibtree.strip() == expected_tree.strip()
def test_edge_dataframe2dict(dframe, expected): assert gruptree.edge_dataframe2dict(dframe) == expected