Skip to content

Commit d017461

Browse files
ryanraaschCDCryanraaschCDC
authored
Merge pull request #52 from cdcent/42-blob-data-mocking
more tests for data sources
2 parents 19e8b60 + dbea4e4 commit d017461

14 files changed

Lines changed: 1211 additions & 471 deletions

‎cfa/scenarios/dataops/datasets/schemas/covid19vax_trends.py‎

Lines changed: 198 additions & 198 deletions
Original file line numberDiff line numberDiff line change
@@ -1,198 +1,198 @@
1-
import random
2-
from typing import Iterable, Optional, Union
3-
4-
import numpy as np
5-
import pandas as pd
6-
import pandera.pandas as pa
7-
from faker import Faker
8-
from pandera import dtypes
9-
from pandera.engines import pandas_engine
10-
11-
fake = Faker()
12-
df_len = 100
13-
14-
15-
@pandas_engine.Engine.register_dtype
16-
@dtypes.immutable
17-
class FixedLenArray(pandas_engine.NpString):
18-
def check(
19-
self,
20-
pandera_dtype: dtypes.DataType,
21-
data_container: Optional[pd.Series] = None,
22-
) -> Union[bool, Iterable[bool]]:
23-
# ensure that the data container's data type is a string,
24-
# using the parent class's check implementation
25-
correct_type = super().check(pandera_dtype)
26-
if not correct_type:
27-
return correct_type
28-
29-
# ensure the filepaths actually exist locally
30-
if data_container is None:
31-
return True
32-
else:
33-
length = len(data_container[0])
34-
return data_container.map(lambda x: len(x) == length)
35-
36-
def __str__(self) -> str:
37-
return str(self.__class__.__name__)
38-
39-
def __repr__(self) -> str:
40-
return f"DataType({self})"
41-
42-
43-
extract_schema = pa.DataFrameSchema(
44-
{
45-
"date": pa.Column(
46-
str,
47-
checks=[
48-
pa.Check.str_matches(
49-
r"^[0-9]{4}-[0-9]{2}-[0-9]{2}(T[0-9]{2}:[0-9]{2}:[0-9]{2}\.[0-9]{3})?$"
50-
)
51-
],
52-
),
53-
"location": pa.Column(
54-
str, checks=[pa.Check.str_matches(r"^[A-Z]{2}$")]
55-
),
56-
"demographic_category": pa.Column(
57-
str, checks=[pa.Check.str_matches(r"^[A-Za-z0-9\_\-\+\<\>]+$")]
58-
),
59-
"census": pa.Column(
60-
int, checks=[pa.Check.in_range(0, 10_000_000_000)], nullable=True
61-
),
62-
"administered_dose1": pa.Column(
63-
int,
64-
checks=[pa.Check.in_range(0, 10_000_000_000)],
65-
nullable=True,
66-
coerce=True,
67-
),
68-
"series_complete_yes": pa.Column(
69-
int,
70-
checks=[pa.Check.in_range(0, 10_000_000_000)],
71-
nullable=True,
72-
coerce=True,
73-
),
74-
"booster_doses": pa.Column(
75-
int,
76-
checks=[pa.Check.in_range(0, 10_000_000_000)],
77-
nullable=True,
78-
coerce=True,
79-
),
80-
"second_booster": pa.Column(
81-
int,
82-
checks=[pa.Check.in_range(0, 10_000_000_000)],
83-
nullable=True,
84-
coerce=True,
85-
),
86-
"administered_dose1_pct_agegroup": pa.Column(
87-
float, checks=[pa.Check.in_range(0, 100)], nullable=True
88-
),
89-
"series_complete_pop_pct_agegroup": pa.Column(
90-
float, checks=[pa.Check.in_range(0, 100)], nullable=True
91-
),
92-
"booster_doses_vax_pct_agegroup": pa.Column(
93-
float, checks=[pa.Check.in_range(0, 100)], nullable=True
94-
),
95-
"second_booster_vax_pct_agegroup": pa.Column(
96-
float, checks=[pa.Check.in_range(0, 100)], nullable=True
97-
),
98-
}
99-
)
100-
101-
load_schema = pa.DataFrameSchema(
102-
{
103-
"date": pa.Column(
104-
str,
105-
checks=[
106-
pa.Check.str_matches(
107-
r"^[0-9]{4}-[0-9]{2}-[0-9]{2}(T[0-9]{2}:[0-9]{2}:[0-9]{2}\.[0-9]{3})?$"
108-
)
109-
],
110-
),
111-
"state": pa.Column(str, checks=[pa.Check.str_matches(r"^[A-Z]{2}$")]),
112-
"age": pa.Column(
113-
str, checks=[pa.Check.isin(["65+", "18-49", "0-17", "50-64"])]
114-
),
115-
"census": pa.Column(
116-
int, checks=[pa.Check.in_range(0, 10_000_000_000)], coerce=True
117-
),
118-
"total": pa.Column(FixedLenArray),
119-
"percentage": pa.Column(FixedLenArray),
120-
"dose": pa.Column(FixedLenArray),
121-
}
122-
)
123-
124-
locations = ["CA", "ME", "NV", "MD"]
125-
ages = ["<5", "5-11_", "12-17_", "18-24_", "25-49_", "50-64_", "65+_"]
126-
tf_ages = ["0-17", "18-49", "50-64", "65+"]
127-
128-
raw_synth_data = pd.DataFrame()
129-
raw_synth_data = raw_synth_data.assign(
130-
date=pd.Series(
131-
str(
132-
fake.date_between_dates(
133-
pd.to_datetime("2021-01-29"), pd.to_datetime("2023-05-10")
134-
)
135-
)
136-
for _ in range(df_len)
137-
),
138-
location=pd.Series(random.choice(locations) for _ in range(df_len)),
139-
demographic_category=pd.Series(
140-
["Ages_" + random.choice(ages) + "yrs" for _ in range(df_len)]
141-
),
142-
census=pd.Series(
143-
fake.unique.random_int(min=200000, max=1000000) for _ in range(df_len)
144-
),
145-
administered_dose1_pct_agegroup=pd.Series(
146-
random.uniform(1.0, 50.0) for _ in range(df_len)
147-
),
148-
series_complete_pop_pct_agegroup=pd.Series(
149-
random.uniform(0.1, 25.0) for _ in range(df_len)
150-
),
151-
booster_doses_vax_pct_agegroup=pd.Series(
152-
random.uniform(0.0, 10.0) for _ in range(df_len)
153-
),
154-
second_booster_vax_pct_agegroup=pd.Series(
155-
random.uniform(0.0, 10.0) for _ in range(df_len)
156-
),
157-
)
158-
raw_synth_data["administered_dose1"] = raw_synth_data["census"] * [
159-
random.uniform(0.1, 0.5) for _ in range(df_len)
160-
]
161-
raw_synth_data["series_complete_yes"] = raw_synth_data[
162-
"administered_dose1"
163-
] * [random.uniform(0.1, 0.5) for _ in range(df_len)]
164-
raw_synth_data["booster_doses"] = raw_synth_data["administered_dose1"] * [
165-
random.uniform(0.1, 0.25) for _ in range(df_len)
166-
]
167-
raw_synth_data["second_booster"] = raw_synth_data["booster_doses"] * [
168-
random.uniform(0.1, 0.5) for _ in range(df_len)
169-
]
170-
171-
tf_synth_data = pd.DataFrame()
172-
tf_synth_data = tf_synth_data.assign(
173-
date=pd.Series(
174-
str(
175-
fake.date_between_dates(
176-
pd.to_datetime("2021-01-29"), pd.to_datetime("2023-05-10")
177-
)
178-
)
179-
for _ in range(df_len)
180-
),
181-
state=pd.Series(random.choice(locations) for _ in range(df_len)),
182-
age=pd.Series([random.choice(tf_ages) for _ in range(df_len)]),
183-
census=pd.Series(
184-
fake.unique.random_int(min=200000, max=1000000) for _ in range(df_len)
185-
),
186-
percentage=pd.Series(
187-
np.array(
188-
[
189-
random.uniform(0.5, 0.99),
190-
random.uniform(0.25, 0.5),
191-
random.uniform(0.1, 0.24),
192-
]
193-
)
194-
for _ in range(df_len)
195-
),
196-
dose=pd.Series([1, 2, 3] for _ in range(df_len)),
197-
)
198-
tf_synth_data["total"] = tf_synth_data["census"] * tf_synth_data["percentage"]
1+
import random
2+
from typing import Iterable, Optional, Union
3+
4+
import numpy as np
5+
import pandas as pd
6+
import pandera.pandas as pa
7+
from faker import Faker
8+
from pandera import dtypes
9+
from pandera.engines import pandas_engine
10+
11+
fake = Faker()
12+
df_len = 100
13+
14+
15+
@pandas_engine.Engine.register_dtype
16+
@dtypes.immutable
17+
class FixedLenArray(pandas_engine.NpString):
18+
def check(
19+
self,
20+
pandera_dtype: dtypes.DataType,
21+
data_container: Optional[pd.Series] = None,
22+
) -> Union[bool, Iterable[bool]]:
23+
# ensure that the data container's data type is a string,
24+
# using the parent class's check implementation
25+
correct_type = super().check(pandera_dtype)
26+
if not correct_type:
27+
return correct_type
28+
29+
# ensure the filepaths actually exist locally
30+
if data_container is None:
31+
return True
32+
else:
33+
length = len(data_container[0])
34+
return data_container.map(lambda x: len(x) == length)
35+
36+
def __str__(self) -> str:
37+
return str(self.__class__.__name__)
38+
39+
def __repr__(self) -> str:
40+
return f"DataType({self})"
41+
42+
43+
extract_schema = pa.DataFrameSchema(
44+
{
45+
"date": pa.Column(
46+
str,
47+
checks=[
48+
pa.Check.str_matches(
49+
r"^[0-9]{4}-[0-9]{2}-[0-9]{2}(T[0-9]{2}:[0-9]{2}:[0-9]{2}\.[0-9]{3})?$"
50+
)
51+
],
52+
),
53+
"location": pa.Column(
54+
str, checks=[pa.Check.str_matches(r"^[A-Z]{2}$")]
55+
),
56+
"demographic_category": pa.Column(
57+
str, checks=[pa.Check.str_matches(r"^[A-Za-z0-9\_\-\+\<\>]+$")]
58+
),
59+
"census": pa.Column(
60+
int, checks=[pa.Check.in_range(0, 10_000_000_000)], nullable=True
61+
),
62+
"administered_dose1": pa.Column(
63+
int,
64+
checks=[pa.Check.in_range(0, 10_000_000_000)],
65+
nullable=True,
66+
coerce=True,
67+
),
68+
"series_complete_yes": pa.Column(
69+
int,
70+
checks=[pa.Check.in_range(0, 10_000_000_000)],
71+
nullable=True,
72+
coerce=True,
73+
),
74+
"booster_doses": pa.Column(
75+
int,
76+
checks=[pa.Check.in_range(0, 10_000_000_000)],
77+
nullable=True,
78+
coerce=True,
79+
),
80+
"second_booster": pa.Column(
81+
int,
82+
checks=[pa.Check.in_range(0, 10_000_000_000)],
83+
nullable=True,
84+
coerce=True,
85+
),
86+
"administered_dose1_pct_agegroup": pa.Column(
87+
float, checks=[pa.Check.in_range(0, 100)], nullable=True
88+
),
89+
"series_complete_pop_pct_agegroup": pa.Column(
90+
float, checks=[pa.Check.in_range(0, 100)], nullable=True
91+
),
92+
"booster_doses_vax_pct_agegroup": pa.Column(
93+
float, checks=[pa.Check.in_range(0, 100)], nullable=True
94+
),
95+
"second_booster_vax_pct_agegroup": pa.Column(
96+
float, checks=[pa.Check.in_range(0, 100)], nullable=True
97+
),
98+
}
99+
)
100+
101+
load_schema = pa.DataFrameSchema(
102+
{
103+
"date": pa.Column(
104+
str,
105+
checks=[
106+
pa.Check.str_matches(
107+
r"^[0-9]{4}-[0-9]{2}-[0-9]{2}(T[0-9]{2}:[0-9]{2}:[0-9]{2}\.[0-9]{3})?$"
108+
)
109+
],
110+
),
111+
"state": pa.Column(str, checks=[pa.Check.str_matches(r"^[A-Z]{2}$")]),
112+
"age": pa.Column(
113+
str, checks=[pa.Check.isin(["65+", "18-49", "0-17", "50-64"])]
114+
),
115+
"census": pa.Column(
116+
int, checks=[pa.Check.in_range(0, 10_000_000_000)], coerce=True
117+
),
118+
"total": pa.Column(FixedLenArray),
119+
"percentage": pa.Column(FixedLenArray),
120+
"dose": pa.Column(FixedLenArray),
121+
}
122+
)
123+
124+
locations = ["CA", "ME", "NV", "MD"]
125+
ages = ["<5", "5-11_", "12-17_", "18-24_", "25-49_", "50-64_", "65+_"]
126+
tf_ages = ["0-17", "18-49", "50-64", "65+"]
127+
128+
raw_synth_data = pd.DataFrame()
129+
raw_synth_data = raw_synth_data.assign(
130+
date=pd.Series(
131+
str(
132+
fake.date_between_dates(
133+
pd.to_datetime("2021-01-29"), pd.to_datetime("2023-05-10")
134+
)
135+
)
136+
for _ in range(df_len)
137+
),
138+
location=pd.Series(random.choice(locations) for _ in range(df_len)),
139+
demographic_category=pd.Series(
140+
["Ages_" + random.choice(ages) + "yrs" for _ in range(df_len)]
141+
),
142+
census=pd.Series(
143+
fake.unique.random_int(min=200000, max=1000000) for _ in range(df_len)
144+
),
145+
administered_dose1_pct_agegroup=pd.Series(
146+
random.uniform(1.0, 50.0) for _ in range(df_len)
147+
),
148+
series_complete_pop_pct_agegroup=pd.Series(
149+
random.uniform(0.1, 25.0) for _ in range(df_len)
150+
),
151+
booster_doses_vax_pct_agegroup=pd.Series(
152+
random.uniform(0.0, 10.0) for _ in range(df_len)
153+
),
154+
second_booster_vax_pct_agegroup=pd.Series(
155+
random.uniform(0.0, 10.0) for _ in range(df_len)
156+
),
157+
)
158+
raw_synth_data["administered_dose1"] = raw_synth_data["census"] * [
159+
random.uniform(0.1, 0.5) for _ in range(df_len)
160+
]
161+
raw_synth_data["series_complete_yes"] = raw_synth_data[
162+
"administered_dose1"
163+
] * [random.uniform(0.1, 0.5) for _ in range(df_len)]
164+
raw_synth_data["booster_doses"] = raw_synth_data["administered_dose1"] * [
165+
random.uniform(0.1, 0.25) for _ in range(df_len)
166+
]
167+
raw_synth_data["second_booster"] = raw_synth_data["booster_doses"] * [
168+
random.uniform(0.1, 0.5) for _ in range(df_len)
169+
]
170+
171+
tf_synth_data = pd.DataFrame()
172+
tf_synth_data = tf_synth_data.assign(
173+
date=pd.Series(
174+
str(
175+
fake.date_between_dates(
176+
pd.to_datetime("2021-01-29"), pd.to_datetime("2023-05-10")
177+
)
178+
)
179+
for _ in range(df_len)
180+
),
181+
state=pd.Series(random.choice(locations) for _ in range(df_len)),
182+
age=pd.Series([random.choice(tf_ages) for _ in range(df_len)]),
183+
census=pd.Series(
184+
fake.unique.random_int(min=200000, max=1000000) for _ in range(df_len)
185+
),
186+
percentage=pd.Series(
187+
np.array(
188+
[
189+
random.uniform(0.5, 0.99),
190+
random.uniform(0.25, 0.5),
191+
random.uniform(0.1, 0.24),
192+
]
193+
)
194+
for _ in range(df_len)
195+
),
196+
dose=pd.Series([1, 2, 3] for _ in range(df_len)),
197+
)
198+
tf_synth_data["total"] = tf_synth_data["census"] * tf_synth_data["percentage"]

0 commit comments

Comments
 (0)