def test_get_attribute():
with pytest.raises(Exception) as e:
atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttribute("emf"))') \
.value()
assert "getAttributeFloat" in str(e.value)
def test_can_call_prodVtx():
ctyp.add_method_type_info(
"xAOD::TruthParticle", "prodVtx",
ctyp.terminal('xAODTruth::TruthVertex', p_depth=1))
atlas_xaod_dataset("file://root.root") \
.Select("lambda e: e.TruthParticles('TruthParticles').Select(lambda t: t.prodVtx().x()).Sum()") \
.value()
def test_dict_output_fail_expansion():
my_old_dict = {1: 'hi'}
with pytest.raises(ValueError):
atlas_xaod_dataset() \
.Select(lambda e: e.EventInfo("EventInfo").runNumber()) \
.Select(lambda e: {'run_number': e, **my_old_dict}) \
.value()
def test_first_object_in_each_event():
# Part of testing that First puts its outer settings in the right place.
# This also tests First on a collection of objects that hasn't been pulled a part
# in a select.
atlas_xaod_dataset() \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").First().pt()/1000.0') \
.value()
def test_first_after_where():
# Part of testing that First puts its outer settings in the right place.
# This also tests First on a collection of objects that hasn't been pulled a part
# in a select.
atlas_xaod_dataset() \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").Where(lambda j: j.pt() > 10).First().pt()') \
.value()
def test_generate_binary_operator_unsupported():
# Make sure an unsupported binary operator triggers an exception
with pytest.raises(Exception) as e:
atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()//2)') \
.value()
assert "FloorDiv" in str(e)
def test_Select_of_2D_array_with_tuple():
# We do not support structured output - so array or array(array), but not array(array, array),
# at least not yet. Make sure error is reasonable.
with pytest.raises(Exception) as e:
atlas_xaod_dataset() \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: (j.pt()/1000.0, j.eta()))') \
.value()
assert "data structures" in str(e.value)
def test_cant_call_double():
msg = ""
try:
atlas_xaod_dataset("file://root.root") \
.Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt().eta()).Sum()") \
.value()
except xAODTranslationError as e:
msg = str(e)
assert 'Unable to call method eta on type double' in msg
def test_per_event_item():
r = atlas_xaod_dataset() \
.Select('lambda e: e.EventInfo("EventInfo").runNumber()') \
.value()
vs = r.QueryVisitor._gc._class_vars
assert 1 == len(vs)
assert "double" == str(vs[0].cpp_type())
def test_generate_unary_not():
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: not (j.pt() > 50.0))') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
_ = find_line_with("!(", lines)
def test_Aggregate_uses_floats_for_float_sum():
r = atlas_xaod_dataset() \
.Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_agg_decl = find_line_with('double agg', lines)
assert l_agg_decl > 0
def test_collection_return():
'Make sure we can set and deal with a returned collection properly'
ctyp.add_method_type_info(
"xAOD::TruthParticle", "vertexes",
ctyp.collection(ctyp.terminal('double'), p_depth=1))
(atlas_xaod_dataset("file://root.root").SelectMany(
lambda e: e.TruthParticles('TruthParticles')).SelectMany(
lambda t: t.vertexes()).value())
def test_Aggregate_not_initial_const_SUM():
r = atlas_xaod_dataset() \
.Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_sets = find_line_numbers_with("/1000", lines)
assert 1 == len(l_sets)
def test_dict_output():
'This is integration testing - making sure the dict to root conversion works'
r = atlas_xaod_dataset() \
.Select(lambda e: e.EventInfo("EventInfo").runNumber()) \
.Select(lambda e: {'run_number': e}) \
.value()
vs = r.QueryVisitor._gc._class_vars
assert 1 == len(vs)
assert "double" == str(vs[0].cpp_type())
def test_First_selects_collection_count():
# Make sure that we have the "First" predicate after if Where's if statement.
r = atlas_xaod_dataset() \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Tracks("InDetTrackParticles")).First().Count()') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
ln = find_line_numbers_with("for", lines)
assert 2 == len(ln)
def test_Aggregate_per_jet_int():
r = atlas_xaod_dataset() \
.Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()).Count()") \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_agg_decl = find_line_with('int agg', lines)
assert l_agg_decl > 0
def test_tree_name():
r = atlas_xaod_dataset() \
.Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \
.AsROOTTTree('junk.root', 'analysis', ['fork']) \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_sets = find_line_numbers_with('tree("analysis")', lines)
assert 1 == len(l_sets)
def test_event_collection_bad_type_arg():
'This is integration testing - making sure the dict to root conversion works'
with pytest.raises(ValueError) as e:
(atlas_xaod_dataset().Select(
lambda e: e.EventInfo(2).runNumber()).Select(lambda e: {
'run_number': e
}).value())
assert "is a string" in str(e)
def test_generate_unary_operations():
ops = ['+', '-']
for o in ops:
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()+({0}1))'.format(o)) \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
_ = find_line_with(f"pt()+({o}(1))", lines)
def test_event_collection_too_many_arg():
'This is integration testing - making sure the dict to root conversion works'
with pytest.raises(ValueError) as e:
(atlas_xaod_dataset().Select(lambda e: e.EventInfo(
"EventInfo", "dork").runNumber()).Select(lambda e: {
'run_number': e
}).value())
assert "only one argument" in str(e)
def test_nested_lambda_argument_name_with_monad():
# Need both the monad and the "e" reused to get this error!
r = atlas_xaod_dataset() \
.Select('lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))') \
.Select('lambda e: e[0].Select(lambda e: e.E())') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_push = find_line_with('push_back', lines)
assert "->E()" in lines[l_push]
def test_electron_and_muon_with_list_qastle():
# See if we can re-create a bug we are seeing with
# Marc's long query.
r = atlas_xaod_dataset(qastle_roundtrip=True) \
.Select('lambda e: [e.Electrons("Electrons"), e.Muons("Muons")]') \
.Select('lambda e: [e[0].Select(lambda ele: ele.E()), e[0].Select(lambda ele: ele.pt()), e[0].Select(lambda ele: ele.phi()), e[0].Select(lambda ele: ele.eta()), e[1].Select(lambda mu: mu.E()), e[1].Select(lambda mu: mu.pt()), e[1].Select(lambda mu: mu.phi()), e[1].Select(lambda mu: mu.eta())]') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
assert find_line_with("->Fill()", lines) != 0
def test_First_Of_Select_After_Where_is_in_right_place():
# Make sure that we have the "First" predicate after if Where's if statement.
r = atlas_xaod_dataset() \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()/1000.0).Where(lambda jpt: jpt > 10.0).First()') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
ln = find_line_with(">10.0", lines)
# Look for the "false" that First uses to remember it has gone by one.
assert find_line_with("false", lines[ln:], throw_if_not_found=False) > 0
def test_generate_binary_operator_pow():
# Make sure the pow operator works correctly - that it doesn't cause a crash in generation.
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()**2)') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l1 = find_line_with("pow(i_obj", lines)
l2 = find_line_with("->pt(), 2)", lines)
assert l1 == l2
def test_get_attribute_float():
# The following statement should be a straight sequence, not an array.
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttributeFloat("emf"))') \
.value()
# Check to see if there mention of push_back anywhere.
lines = get_lines_of_code(r)
print_lines(lines)
l_attribute = find_line_with("getAttribute<", lines)
assert 'getAttribute<float>("emf")' in lines[l_attribute]
def test_generate_binary_operators():
# Make sure the binary operators work correctly - that they don't cause a crash in generation.
ops = ['+', '-', '*', '/', '%']
for o in ops:
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt(){0}1)'.format(o)) \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
_ = find_line_with(f"pt(){o}1", lines)
def test_builtin_sin_function_no_math_import():
# The following statement should be a straight sequence, not an array.
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: sin(j.pt()))') \
.value()
# Check to see if there mention of push_back anywhere.
lines = get_lines_of_code(r)
print_lines(lines)
l_abs = find_line_with("std::sin", lines)
assert "->pt()" in lines[l_abs]
def test_ifexpr():
r = atlas_xaod_dataset(qastle_roundtrip=True) \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: 1.0 if j.pt() > 10.0 else 2.0)') \
.value()
# Make sure that a test around 10.0 occurs.
lines = get_lines_of_code(r)
print_lines(lines)
lines = [ln for ln in lines if '10.0' in ln]
assert len(lines) == 1
assert 'if ' in lines[0]
def test_SelectMany_of_tuple_is_not_array():
# The following statement should be a straight sequence, not an array.
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: (j.pt()/1000.0, j.eta()))') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
assert 0 == ["push_back" in ln for ln in lines].count(True)
l_push_back = find_line_with("Fill()", lines)
active_blocks = find_open_blocks(lines[:l_push_back])
assert 1 == ["for" in a for a in active_blocks].count(True)
def test_result_awkward():
# The following statement should be a straight sequence, not an array.
r = atlas_xaod_dataset() \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \
.Select("lambda j: j.pt()") \
.value()
# Make sure that the tree Fill is at the same level as the _JetPts2 getting set.
lines = get_lines_of_code(r)
print_lines(lines)
l_jetpt = find_line_with("_col1", lines)
assert "Fill()" in lines[l_jetpt + 1]
