Skip to content

Commit 07a8579

Browse files
BenGalewskyBenGalewsky
committed
Add unit tests for databinder
1 parent 3e609b2 commit 07a8579

File tree

2 files changed

+134
-208
lines changed

2 files changed

+134
-208
lines changed

‎servicex/databinder_models.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
2828
from enum import Enum
2929
from typing import Union, Optional, Callable, List
30-
from pydantic import BaseModel, Field, AnyUrl, root_validator, constr, validator
30+
from pydantic import BaseModel, Field, root_validator, constr, validator
3131

3232
from servicex.dataset_identifier import RucioDatasetIdentifier, FileListDataset
3333
from servicex.func_adl import func_adl_dataset
@@ -95,7 +95,7 @@ class DeliveryEnum(str, Enum):
9595
OutputFormat: Union[OutputFormatEnum,
9696
constr(regex='^(parquet|root-file)$')] = Field(default=OutputFormatEnum.root) # NOQA F722
9797

98-
Delivery: Union[DeliveryEnum, constr(regex='^(LocalCache|SignedURLs)$')] # NOQA F722
98+
Delivery: Union[DeliveryEnum, constr(regex='^(LocalCache|SignedURLs)$')] = Field(default=DeliveryEnum.LocalCache) # NOQA F722
9999

100100

101101
class Definition(BaseModel):

‎tests/test_databinder.py

Lines changed: 132 additions & 206 deletions
Original file line numberDiff line numberDiff line change
@@ -1,226 +1,152 @@
1-
from unittest.mock import patch
21
import pytest
3-
from pathlib import Path
4-
import yaml
2+
from pydantic import ValidationError
53

6-
from servicex.models import TransformedResults
7-
# from servicex.dataset import Dataset
4+
from servicex import ServiceXSpec, FileListDataset, RucioDatasetIdentifier
85

9-
from servicex.databinder.databinder_configuration \
10-
import load_databinder_config, _set_default_values, \
11-
_support_old_option_names
12-
from servicex.databinder.databinder_requests import DataBinderRequests
13-
from servicex.databinder.databinder_outputs import OutputHandler
14-
from servicex.databinder.databinder_deliver import DataBinderDeliver
15-
from servicex.databinder.databinder import DataBinder
166

17-
18-
def test_load_config():
19-
config = {
20-
"General":
21-
{
7+
def basic_spec(samples=None):
8+
return {
9+
"General": {
2210
"ServiceX": "servicex",
2311
"Codegen": "python",
24-
"Delivery": "ObjectStore"
2512
},
26-
"Sample":
27-
[
28-
{
29-
"Name": "sampleA",
30-
"RucioDID": "user.kchoi:sampleA",
31-
"Function": "DEF_a"
32-
},
33-
{
34-
"Name": "sampleB",
35-
"XRootDFiles": "root://a.root",
36-
"Columns": "el_pt",
37-
"Codegen": "uproot"
38-
}
39-
],
40-
"Definition":
41-
{
42-
"DEF_a": "a"
43-
}
44-
}
45-
new_config = load_databinder_config(config)
46-
assert new_config["Sample"][0]["Function"] == "a"
47-
48-
with open("temp_databinder.yaml", "w") as f:
49-
yaml.dump(new_config, f, default_flow_style=False)
50-
new_config = load_databinder_config("temp_databinder.yaml")
51-
assert new_config["Sample"][0]["Function"] == "a"
52-
53-
54-
def test_config_default_value():
55-
config = {
56-
"General": {
57-
"IgnoreLocalCache": True
58-
},
59-
"Sample": [
60-
{
61-
"Name": "sampleA",
62-
"RucioDID": "user.kchoi:sampleA",
63-
"Function": "DEF_a"
64-
}
65-
]
13+
"Sample": samples
14+
or [{"Name": "sampleA", "XRootDFiles": "root://a.root", "Function": "a"}],
6615
}
67-
new_config = _set_default_values(config)
68-
assert new_config["General"]["Delivery"] == "localpath"
6916

7017

71-
def test_config_old_option_names():
18+
def test_load_config():
7219
config = {
7320
"General": {
74-
"ServiceXName": "servicex",
75-
"Transformer": "uproot",
76-
"IgnoreServiceXCache": True
77-
},
78-
"Sample": [
79-
{
80-
"Name": "sampleA",
81-
"RucioDID": "user.kchoi:sampleA",
82-
"Transformer": "python",
83-
"Function": "DEF_a",
84-
"IgnoreServiceXCache": True
85-
}
86-
]
87-
}
88-
new_config = _support_old_option_names(config)
89-
assert new_config["General"]["ServiceX"] == "servicex"
90-
91-
92-
@patch('servicex.databinder.databinder_requests.DataBinderRequests._get_client')
93-
def test_requests_python_transformer(_get_client):
94-
config = {
95-
"General":
96-
{
9721
"ServiceX": "servicex",
98-
"OutputFormat": "root",
99-
"Delivery": "objectstore"
100-
},
101-
"Sample":
102-
[
103-
{
104-
"Name": "sampleA",
105-
"RucioDID": "user.kchoi:sampleA",
106-
"Codegen": "python",
107-
"Function": "DEF_a",
108-
"NFiles": "5",
109-
"IgnoreLocalCache": "False"
110-
},
111-
{
112-
"Name": "sampleB",
113-
"XRootDFiles": "root://a.root",
114-
"Function": "DEF_a",
115-
"Codegen": "python",
116-
"IgnoreLocalCache": "False"
117-
}
118-
]
119-
}
120-
reqs = DataBinderRequests(config).get_requests()
121-
assert reqs[1]["sample_name"] == "sampleB"
122-
assert len(reqs) == 2
123-
assert len(reqs[0].keys()) == 3
124-
# assert isinstance(reqs[0]["ds_query"], Dataset)
125-
126-
127-
def test_output_handler():
128-
config = {
129-
"General":
130-
{
131-
"OutputDirectory": "./temp_dir2",
132-
"Delivery": "objectstore",
133-
"OutFilesetName": "out_dict"
134-
},
135-
"Sample":
136-
[
137-
{
138-
"Name": "sampleA",
139-
},
140-
{
141-
"Name": "sampleB",
142-
}
143-
]
144-
}
145-
result = TransformedResults
146-
result.title = "sampleA"
147-
result.signed_url_list = ["a", "b"]
148-
result.file_list = ["c", "d"]
149-
150-
out = OutputHandler(config)
151-
152-
out.update_out_dict("objectstore", result)
153-
assert len(out.out_dict['samples']['sampleA']) == 2
154-
155-
out.write_out_dict()
156-
with open(Path("temp_dir2/out_dict.yml"), "r") as f:
157-
out_yaml = yaml.safe_load(f)
158-
assert len(out_yaml["samples"]["sampleA"]) == 2
159-
160-
161-
@pytest.mark.asyncio
162-
@patch('servicex.databinder.databinder_requests.DataBinderRequests._get_client')
163-
async def test_deliver(_get_client):
164-
config = {
165-
"General":
166-
{
167-
"ServiceX": "servicex",
168-
"OutputDirectory": "./temp_dir2",
169-
"Delivery": "objectstore",
170-
"OutFilesetName": "out_dict",
171-
"OutputFormat": "root"
22+
"Codegen": "python",
23+
"Delivery": "LocalCache",
17224
},
173-
"Sample":
174-
[
175-
{
176-
"Name": "sampleA",
177-
"RucioDID": "user.kchoi:sampleA",
178-
"Codegen": "python",
179-
"Function": "DEF_a",
180-
"NFiles": "5",
181-
"IgnoreLocalCache": "False"
182-
},
25+
"Sample": [
26+
{"Name": "sampleA", "RucioDID": "user.kchoi:sampleA", "Function": "DEF_a"},
18327
{
18428
"Name": "sampleB",
18529
"XRootDFiles": "root://a.root",
30+
"Columns": "el_pt",
31+
"Codegen": "uproot",
18632
"Function": "DEF_a",
187-
"Codegen": "python",
188-
"IgnoreLocalCache": "False"
189-
}
190-
]
191-
}
192-
deliv = DataBinderDeliver(config)
193-
194-
assert deliv._requests[0]['sample_name'] == "sampleA"
195-
# assert deliv.deliver_and_copy()
196-
197-
198-
@patch('servicex.databinder.databinder_requests.DataBinderRequests._get_client')
199-
def test_databinder(_get_client):
200-
config = {
201-
"General":
202-
{
203-
"ServiceX": "servicex",
204-
"OutputDirectory": "./temp_dir2",
205-
"OutFilesetName": "out_dict",
206-
},
207-
"Sample":
208-
[
209-
{
210-
"Name": "sampleA",
211-
"RucioDID": "user.kchoi:sampleA",
212-
"Codegen": "python",
213-
"Function": "DEF_a",
214-
"NFiles": "5",
21533
},
216-
{
217-
"Name": "sampleB",
218-
"XRootDFiles": "root://a.root",
219-
"Function": "DEF_a",
220-
"Codegen": "python",
221-
}
222-
]
34+
],
35+
"Definition": {"DEF_a": "a"},
22336
}
224-
sx_db = DataBinder(config)
225-
226-
assert sx_db._config["General"]["OutputFormat"] == "root"
37+
new_config = ServiceXSpec.parse_obj(config)
38+
assert new_config.Sample[0].Function == "a"
39+
40+
41+
def test_single_root_file():
42+
43+
spec = ServiceXSpec.parse_obj(
44+
basic_spec(
45+
samples=[
46+
{
47+
"Name": "sampleA",
48+
"XRootDFiles": "root://eospublic.cern.ch//file1.root",
49+
"Function": "a",
50+
}
51+
]
52+
)
53+
)
54+
55+
assert isinstance(spec.Sample[0].dataset_identifier, FileListDataset)
56+
assert spec.Sample[0].dataset_identifier.files == [
57+
"root://eospublic.cern.ch//file1.root"
58+
]
59+
60+
61+
def test_list_of_root_files():
62+
spec = ServiceXSpec.parse_obj(
63+
basic_spec(
64+
samples=[
65+
{
66+
"Name": "sampleA",
67+
"XRootDFiles": [
68+
"root://eospublic.cern.ch//file1.root",
69+
"root://eospublic.cern.ch//file2.root",
70+
],
71+
"Function": "a",
72+
}
73+
]
74+
)
75+
)
76+
77+
assert isinstance(spec.Sample[0].dataset_identifier, FileListDataset)
78+
assert spec.Sample[0].dataset_identifier.files == [
79+
"root://eospublic.cern.ch//file1.root",
80+
"root://eospublic.cern.ch//file2.root",
81+
]
82+
83+
84+
def test_rucio_did():
85+
spec = ServiceXSpec.parse_obj(
86+
basic_spec(
87+
samples=[
88+
{
89+
"Name": "sampleA",
90+
"RucioDID": "user.ivukotic:user.ivukotic.single_top_tW__nominal",
91+
"Function": "a",
92+
}
93+
]
94+
)
95+
)
96+
97+
assert isinstance(spec.Sample[0].dataset_identifier, RucioDatasetIdentifier)
98+
assert (
99+
spec.Sample[0].dataset_identifier.did
100+
== "rucio://user.ivukotic:user.ivukotic.single_top_tW__nominal"
101+
)
102+
103+
104+
def test_rucio_did_numfiles():
105+
spec = ServiceXSpec.parse_obj(
106+
basic_spec(
107+
samples=[
108+
{
109+
"Name": "sampleA",
110+
"RucioDID": "user.ivukotic:user.ivukotic.single_top_tW__nominal",
111+
"NFiles": 10,
112+
"Function": "a",
113+
}
114+
]
115+
)
116+
)
117+
118+
assert isinstance(spec.Sample[0].dataset_identifier, RucioDatasetIdentifier)
119+
assert (
120+
spec.Sample[0].dataset_identifier.did
121+
== "rucio://user.ivukotic:user.ivukotic.single_top_tW__nominal?files=10"
122+
)
123+
124+
125+
def test_invalid_dataset_identifier():
126+
with pytest.raises(ValidationError):
127+
ServiceXSpec.parse_obj(
128+
basic_spec(
129+
samples=[
130+
{
131+
"Name": "sampleA",
132+
"RucioDID": "user.ivukotic:user.ivukotic.single_top_tW__nominal",
133+
"XRootDFiles": "root://eospublic.cern.ch//file1.root",
134+
"NFiles": 10,
135+
"Function": "a",
136+
}
137+
]
138+
)
139+
)
140+
141+
with pytest.raises(ValidationError):
142+
ServiceXSpec.parse_obj(
143+
basic_spec(
144+
samples=[
145+
{
146+
"Name": "sampleA",
147+
"NFiles": 10,
148+
"Function": "a",
149+
}
150+
]
151+
)
152+
)

0 commit comments

Comments
 (0)