|
34 | 34 | import sys
|
35 | 35 | import numpy as np
|
36 | 36 | from servicex_analysis_utils.materialization import to_awk
|
| 37 | +import types |
37 | 38 |
|
38 | 39 |
|
39 | 40 | @pytest.fixture
|
@@ -69,8 +70,10 @@ def test_to_awk(build_test_samples):
|
69 | 70 |
|
70 | 71 | #Collecting all samples
|
71 | 72 | assert list(result.keys())==["Test-Sample1", "Test-Sample2"]
|
72 |
| - arr1 = ak.concatenate(list(result["Test-Sample1"])) # Materialize the generator from uproot.iterate |
73 |
| - arr2 = ak.concatenate(list(result["Test-Sample2"])) |
| 73 | + arr1 = result["Test-Sample1"] |
| 74 | + arr2 = result["Test-Sample2"] |
| 75 | + |
| 76 | + |
74 | 77 |
|
75 | 78 | #Collecting all branches
|
76 | 79 | assert ak.fields(arr1) == ['branch1', 'branch2']
|
@@ -114,6 +117,25 @@ def test_to_awk_dask(build_test_samples):
|
114 | 117 | assert ak.all(arr1['branch2'].compute() == ak.from_numpy(np.zeros(100)))
|
115 | 118 | assert ak.all(arr2['branch1'].compute() == ak.from_numpy(np.ones(10)))
|
116 | 119 |
|
| 120 | +def test_to_awk_delayed_and_kwargs(build_test_samples): |
| 121 | + sx_dict = build_test_samples |
| 122 | + result_delay = to_awk(sx_dict, iterator=True, expressions="branch1") #return iterable + selection kwarg |
| 123 | + |
| 124 | + #Checking iterator return type |
| 125 | + assert isinstance(result_delay["Test-Sample1"], types.GeneratorType) |
| 126 | + assert isinstance(result_delay["Test-Sample2"], types.GeneratorType) |
| 127 | + |
| 128 | + |
| 129 | + arr1 = ak.concatenate(list(result_delay["Test-Sample1"])) # Materialize the generator from uproot.iterate |
| 130 | + arr2 = ak.concatenate(list(result_delay["Test-Sample2"])) |
| 131 | + |
| 132 | + #Checking materialization |
| 133 | + assert isinstance(arr1, ak.Array), "to_awk(dask=True) does not produce an ak.Array instance" |
| 134 | + assert isinstance(arr2, ak.Array), "to_awk(dask=True) does not produce an ak.Array instance" |
| 135 | + |
| 136 | + #Checking only 1 branch selected |
| 137 | + assert ak.fields(arr1) == ['branch1'] |
| 138 | + assert ak.fields(arr2) == ['branch1'] |
117 | 139 |
|
118 | 140 | def test_unsupported_file_format():
|
119 | 141 | fake_paths = {"fake-Sample": ["invalid_file.txt"]}
|
|
0 commit comments