예제 #1
0
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
예제 #2
0
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"])        
예제 #3
0
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"])
예제 #4
0
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
예제 #5
0
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
예제 #6
0
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
예제 #7
0
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)
예제 #8
0
def test_Dataset_resolveWildcardBranch(selector, inbranches, expectBranches):
    newDataset = Dataset("someName") 
    newDataset.filesAdded = True
    newDataset.branches = inbranches
    
    assert expectBranches == newDataset._resolveWildcardBranch(selector)