Skip to content
This repository was archived by the owner on Sep 2, 2025. It is now read-only.

Commit 03f68e3

Browse files
syou6162syou6162
committed
Add tests for data profile scan
1 parent 7f80a97 commit 03f68e3

File tree

1 file changed

+257
-0
lines changed

1 file changed

+257
-0
lines changed

‎tests/functional/test_data_profile_scan.py‎

Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
import pytest
2+
from unittest.mock import patch
3+
from dbt.adapters.bigquery.relation import BigQueryRelation
4+
from dbt.tests.util import run_dbt, get_connection, relation_from_name
5+
6+
SCAN_LOCATION = "us-central1"
7+
SCAN_ID = "bigquery_data_profile_scan_test"
8+
MODEL_NAME = "test_model"
9+
10+
ORIGINAL_LABELS = {
11+
"my_label_key": "my_label_value",
12+
}
13+
14+
PROFILE_SCAN_LABELS = [
15+
"dataplex-dp-published-scan",
16+
"dataplex-dp-published-project",
17+
"dataplex-dp-published-location",
18+
]
19+
20+
SQL_CONTENT = """
21+
{{
22+
config(
23+
materialized="table"
24+
)
25+
}}
26+
select 20 as id, cast('2020-01-01 01:00:00' as datetime) as date_hour union all
27+
select 40 as id, cast('2020-01-01 02:00:00' as datetime) as date_hour
28+
"""
29+
30+
YAML_CONTENT = f"""version: 2
31+
models:
32+
- name: {MODEL_NAME}
33+
"""
34+
35+
YAML_CONTENT_WITH_PROFILE_SCAN_SETTING = f"""version: 2
36+
models:
37+
- name: {MODEL_NAME}
38+
config:
39+
data_profile_scan:
40+
location: us-central1
41+
scan_id: yasuhisa-test4
42+
sampling_percent: 10
43+
row_filter: "TRUE"
44+
cron: "CRON_TZ=Asia/New_York 0 9 * * *"
45+
"""
46+
47+
48+
class TestDataProfileScanWithProjectProfileScanSetting:
49+
@pytest.fixture(scope="class")
50+
def project_config_update(self):
51+
return {
52+
"models": {
53+
"+labels": ORIGINAL_LABELS,
54+
"+data_profile_scan": {
55+
"location": SCAN_LOCATION,
56+
"scan_id": SCAN_ID,
57+
"sampling_percent": 10,
58+
"row_filter": "TRUE",
59+
},
60+
},
61+
}
62+
63+
@pytest.fixture(scope="class")
64+
def models(self):
65+
return {
66+
f"{MODEL_NAME}.sql": SQL_CONTENT,
67+
f"{MODEL_NAME}.yml": YAML_CONTENT,
68+
}
69+
70+
def test_create_data_profile_scan(self, project):
71+
with patch(
72+
"dbt.adapters.bigquery.impl.dataplex_v1.DataScanServiceClient"
73+
) as MockDataScanClient:
74+
mock_data_scan_client = MockDataScanClient.return_value
75+
76+
results = run_dbt()
77+
assert len(results) == 1
78+
79+
mock_data_scan_client.create_data_scan.assert_called_once()
80+
mock_data_scan_client.run_data_scan.assert_called_once()
81+
82+
relation: BigQueryRelation = relation_from_name(project.adapter, MODEL_NAME)
83+
adapter = project.adapter
84+
with get_connection(project.adapter) as conn:
85+
table = conn.handle.get_table(
86+
adapter.connections.get_bq_table(
87+
relation.database, relation.schema, relation.table
88+
)
89+
)
90+
labels_to_be_created = PROFILE_SCAN_LABELS + list(ORIGINAL_LABELS.keys())
91+
assert set(table.labels.keys()) == set(labels_to_be_created)
92+
93+
94+
class TestDataProfileScanWithProjectProfileScanSettingAndCron:
95+
@pytest.fixture(scope="class")
96+
def project_config_update(self):
97+
return {
98+
"models": {
99+
"+labels": ORIGINAL_LABELS,
100+
"+data_profile_scan": {
101+
"location": SCAN_LOCATION,
102+
"scan_id": SCAN_ID,
103+
"sampling_percent": 10,
104+
"row_filter": "TRUE",
105+
"cron": "CRON_TZ=Asia/New_York 0 9 * * *",
106+
},
107+
},
108+
}
109+
110+
@pytest.fixture(scope="class")
111+
def models(self):
112+
return {
113+
f"{MODEL_NAME}.sql": SQL_CONTENT,
114+
f"{MODEL_NAME}.yml": YAML_CONTENT,
115+
}
116+
117+
def test_create_data_profile_scan(self, project):
118+
with patch(
119+
"dbt.adapters.bigquery.impl.dataplex_v1.DataScanServiceClient"
120+
) as MockDataScanClient:
121+
mock_data_scan_client = MockDataScanClient.return_value
122+
123+
results = run_dbt()
124+
assert len(results) == 1
125+
126+
mock_data_scan_client.create_data_scan.assert_called_once()
127+
mock_data_scan_client.run_data_scan.assert_not_called()
128+
129+
relation: BigQueryRelation = relation_from_name(project.adapter, MODEL_NAME)
130+
adapter = project.adapter
131+
with get_connection(project.adapter) as conn:
132+
table = conn.handle.get_table(
133+
adapter.connections.get_bq_table(
134+
relation.database, relation.schema, relation.table
135+
)
136+
)
137+
labels_to_be_created = PROFILE_SCAN_LABELS + list(ORIGINAL_LABELS.keys())
138+
assert set(table.labels.keys()) == set(labels_to_be_created)
139+
140+
141+
class TestDataProfileScanWithModelProfileScanSetting:
142+
@pytest.fixture(scope="class")
143+
def models(self):
144+
sql_content = f"""
145+
{{{{
146+
config(
147+
materialized="table",
148+
labels={ORIGINAL_LABELS},
149+
)
150+
}}}}
151+
select 20 as id, cast('2020-01-01 01:00:00' as datetime) as date_hour union all
152+
select 40 as id, cast('2020-01-01 02:00:00' as datetime) as date_hour
153+
"""
154+
155+
return {
156+
f"{MODEL_NAME}.sql": sql_content,
157+
f"{MODEL_NAME}.yml": YAML_CONTENT_WITH_PROFILE_SCAN_SETTING,
158+
}
159+
160+
def test_create_data_profile_scan(self, project):
161+
with patch(
162+
"dbt.adapters.bigquery.impl.dataplex_v1.DataScanServiceClient"
163+
) as MockDataScanClient:
164+
mock_data_scan_client = MockDataScanClient.return_value
165+
166+
results = run_dbt()
167+
assert len(results) == 1
168+
169+
mock_data_scan_client.create_data_scan.assert_called_once()
170+
mock_data_scan_client.run_data_scan.assert_not_called()
171+
172+
relation: BigQueryRelation = relation_from_name(project.adapter, MODEL_NAME)
173+
adapter = project.adapter
174+
with get_connection(project.adapter) as conn:
175+
table = conn.handle.get_table(
176+
adapter.connections.get_bq_table(
177+
relation.database, relation.schema, relation.table
178+
)
179+
)
180+
labels_to_be_created = PROFILE_SCAN_LABELS + list(ORIGINAL_LABELS.keys())
181+
assert set(table.labels.keys()) == set(labels_to_be_created)
182+
183+
184+
class TestDataProfileScanWithoutProfileScanSetting:
185+
@pytest.fixture(scope="class")
186+
def models(self):
187+
return {
188+
f"{MODEL_NAME}.sql": SQL_CONTENT,
189+
f"{MODEL_NAME}.yml": YAML_CONTENT,
190+
}
191+
192+
def test_create_data_profile_scan(self, project):
193+
with patch(
194+
"dbt.adapters.bigquery.impl.dataplex_v1.DataScanServiceClient"
195+
) as MockDataScanClient:
196+
mock_data_scan_client = MockDataScanClient.return_value
197+
198+
results = run_dbt()
199+
assert len(results) == 1
200+
201+
mock_data_scan_client.create_data_scan.assert_not_called()
202+
mock_data_scan_client.run_data_scan.assert_not_called()
203+
204+
relation: BigQueryRelation = relation_from_name(project.adapter, MODEL_NAME)
205+
adapter = project.adapter
206+
with get_connection(project.adapter) as conn:
207+
table = conn.handle.get_table(
208+
adapter.connections.get_bq_table(
209+
relation.database, relation.schema, relation.table
210+
)
211+
)
212+
labels_to_be_created = []
213+
assert set(table.labels.keys()) == set(labels_to_be_created)
214+
215+
216+
class TestDataProfileScanDisabledProfileScanSetting:
217+
@pytest.fixture(scope="class")
218+
def project_config_update(self):
219+
return {
220+
"models": {
221+
"+data_profile_scan": {
222+
"location": SCAN_LOCATION,
223+
"scan_id": SCAN_ID,
224+
"enabled": False,
225+
},
226+
},
227+
}
228+
229+
@pytest.fixture(scope="class")
230+
def models(self):
231+
return {
232+
f"{MODEL_NAME}.sql": SQL_CONTENT,
233+
f"{MODEL_NAME}.yml": YAML_CONTENT,
234+
}
235+
236+
def test_create_data_profile_scan(self, project):
237+
with patch(
238+
"dbt.adapters.bigquery.impl.dataplex_v1.DataScanServiceClient"
239+
) as MockDataScanClient:
240+
mock_data_scan_client = MockDataScanClient.return_value
241+
242+
results = run_dbt()
243+
assert len(results) == 1
244+
245+
mock_data_scan_client.create_data_scan.assert_not_called()
246+
mock_data_scan_client.run_data_scan.assert_not_called()
247+
248+
relation: BigQueryRelation = relation_from_name(project.adapter, MODEL_NAME)
249+
adapter = project.adapter
250+
with get_connection(project.adapter) as conn:
251+
table = conn.handle.get_table(
252+
adapter.connections.get_bq_table(
253+
relation.database, relation.schema, relation.table
254+
)
255+
)
256+
labels_to_be_created = []
257+
assert set(table.labels.keys()) == set(labels_to_be_created)

0 commit comments

Comments
 (0)