def test_Dataset_setOutputBranches(mockTree):
newDataset = Dataset("someName")
inBrList = [br.decode(newDataset.encoding) for br in mockTree.keys()]
newDataset.filesAdded = True
newDataset.branches = inBrList
expectedBranches = inBrList[0:1]
newDataset.setOutputBranches(expectedBranches)
assert expectedBranches == newDataset.outputBranches and newDataset.outputBranchesSet
def test_Dataset_addFiles_exception(mockTree, mocker):
newDataset = Dataset("someName")
def openROOTFile(inputfile):
return {newDataset.treeName : mockTree}
mocker.patch("uproot.open", new=openROOTFile)
newDataset.filesAdded = True
newDataset.branches = ["branch5"]
with pytest.raises(RuntimeError):
newDataset.addFiles(["someOtherFile.root"])
def test_Dataset_setOutputBranches_exception(mockTree):
newDataset = Dataset("someName")
inBrList = [br.decode(newDataset.encoding) for br in mockTree.keys()]
expectedBranches = inBrList[0:1]
with pytest.raises(RuntimeError):
newDataset.setOutputBranches(expectedBranches)
newDataset.filesAdded = True
newDataset.branches = inBrList
with pytest.raises(KeyError):
newDataset.setOutputBranches(["noPresentBranch"])
def test_Dataset_cleanBranchList(mockTree):
newDataset = Dataset("someName")
inBrList = [br.decode(newDataset.encoding) for br in mockTree.keys()]
newDataset.filesAdded = True
newDataset.branches = inBrList
expectedBranches = inBrList[0:1]
newDataset.setOutputBranches(expectedBranches)
assert expectedBranches == newDataset.outputBranches
newDataset.cleanBranchList([expectedBranches[0]])
expectedBranches.remove(expectedBranches[0])
assert expectedBranches == newDataset.outputBranches
def test_Dataset_setOutputBranches_wildcard_contains(selector, expectBranches, count, mocker, mockTree):
newDataset = Dataset("someName")
inBrList = ["var_AAA_1","var_AAA_2","var_BBB_1","var_BBB_2","var_CCC_1"]
newDataset.filesAdded = True
newDataset.branches = inBrList
expectedBranches = expectBranches
mocker.spy(newDataset, "_resolveWildcardBranch")
newDataset.setOutputBranches(selector)
assert newDataset._resolveWildcardBranch.call_count == count
assert sorted(expectedBranches) == sorted(newDataset.outputBranches) and newDataset.outputBranchesSet
def test_Dataset_setOutputBranches_wildcard_all(mocker, mockTree):
newDataset = Dataset("someName")
inBrList = [br.decode(newDataset.encoding) for br in mockTree.keys()]
newDataset.filesAdded = True
newDataset.branches = inBrList
expectedBranches = inBrList
mocker.spy(newDataset, "_resolveWildcardBranch")
newDataset.setOutputBranches("*")
assert newDataset._resolveWildcardBranch.call_count == 0
assert sorted(expectedBranches) == sorted(newDataset.outputBranches) and newDataset.outputBranchesSet
def test_Dataset_process(mockTree, mocker):
newDataset = Dataset("someName")
mockTree_1 = copy.deepcopy(mockTree)
mockTree_2 = copy.deepcopy(mockTree)
mockTree_1.dataframe.update(pd.DataFrame({'branch2': list(range(2,12))[::-1]}))
mockTree_1.setDF()
mockTree_2.dataframe.update(pd.DataFrame({'branch1': (list(range(0,10)))[::-1]}))
mockTree_2.setDF()
newDataset.filesAdded = True
newDataset.files = ["file1.root", "file2.root"]
newDataset.branches = ["branch1","branch2","branch3"]
newDataset.outputBranchesSet = True
newDataset.outputBranches = ["branch1", "branch3"]
def openROOTFile(*args, **kwargs):
mm = mocker.MagicMock()
inputfile = args[0]
if inputfile == "file1.root":
mm.__enter__ = mocker.Mock(return_value =
{newDataset.treeName : copy.deepcopy(mockTree_1)}
)
else:
mm.__enter__ = mocker.Mock(return_value =
{newDataset.treeName : copy.deepcopy(mockTree_2)}
)
return mm
m = mocker.MagicMock() #This mocker, mocks the open call
m.side_effect = openROOTFile #Returns a mocker to deal with the with statement
mocker.patch("uproot.open", m, create=True)
newDataset.selection = "branch1 >= 7 and branch2 >=2"
mockTree_1_df = mockTree_1.dataframe
mockTree_2_df = mockTree_2.dataframe
mockTree_1_df = mockTree_1_df.query("branch1 >= 7 and branch2 >=2")
mockTree_2_df = mockTree_2_df.query("branch1 >= 7 and branch2 >=2")
expected = pd.concat([mockTree_1_df, mockTree_2_df])
expected.drop(columns=["branch2"], inplace=True)
outputDF = newDataset.process(skipOutput = True)
assert outputDF.equals(expected)
def test_Dataset_resolveWildcardBranch(selector, inbranches, expectBranches):
newDataset = Dataset("someName")
newDataset.filesAdded = True
newDataset.branches = inbranches
assert expectBranches == newDataset._resolveWildcardBranch(selector)