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
