Skip to content

Commit

Permalink
add case for incremental model
Browse files Browse the repository at this point in the history
  • Loading branch information
syou6162 committed Nov 4, 2024
1 parent 8a99bfe commit cd6965b
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,7 @@
{% do apply_grants(target_relation, grant_config, should_revoke) %}

{% do persist_docs(target_relation, model) %}
{% do adapter.create_or_update_data_profile_scan(model) %}

{%- if tmp_relation_exists -%}
{{ adapter.drop_relation(tmp_relation) }}
Expand Down
74 changes: 74 additions & 0 deletions tests/functional/test_data_profile_scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,26 @@
cron: "CRON_TZ=Asia/New_York 0 9 * * *"
"""

INCREMENTAL_MODEL_CONTENT = """
{{
config(
materialized="incremental",
)
}}
{% if not is_incremental() %}
select 10 as id, cast('2020-01-01 01:00:00' as datetime) as date_hour union all
select 30 as id, cast('2020-01-01 02:00:00' as datetime) as date_hour
{% else %}
select 20 as id, cast('2020-01-01 01:00:00' as datetime) as date_hour union all
select 40 as id, cast('2020-01-01 02:00:00' as datetime) as date_hour
{% endif %}
"""


class TestDataProfileScanWithProjectProfileScanSetting:
@pytest.fixture(scope="class")
Expand Down Expand Up @@ -362,3 +382,57 @@ def test_create_data_profile_scan(self, project):
)
labels_to_be_created = list(ORIGINAL_LABELS.keys())
assert set(table.labels.keys()) == set(labels_to_be_created)


class TestDataProfileScanWithIncrementalModel:
@pytest.fixture(scope="class")
def project_config_update(self):
return {
"models": {
"+labels": ORIGINAL_LABELS,
"+data_profile_scan": {
"location": SCAN_LOCATION,
"scan_id": SCAN_ID,
"sampling_percent": 10,
"row_filter": "TRUE",
},
},
}

@pytest.fixture(scope="class")
def models(self):
return {
f"{MODEL_NAME}.sql": INCREMENTAL_MODEL_CONTENT,
f"{MODEL_NAME}.yml": YAML_CONTENT,
}

def test_create_data_profile_scan(self, project):
with patch("google.cloud.dataplex_v1.DataScanServiceClient") as MockDataScanClient:
mock_data_scan_client = MockDataScanClient.return_value

results = run_dbt()
assert len(results) == 1

mock_data_scan_client.create_data_scan.assert_called_once()
mock_data_scan_client.run_data_scan.assert_called_once()

def list_data_scans_mock(parent):
mock_scan = MagicMock()
mock_scan.name = SCAN_ID
return [mock_scan]

mock_data_scan_client.list_data_scans.side_effect = list_data_scans_mock

results = run_dbt()
assert len(results) == 1
mock_data_scan_client.update_data_scan.assert_called_once()

relation: BigQueryRelation = relation_from_name(project.adapter, MODEL_NAME)
with get_connection(project.adapter) as conn:
table = conn.handle.get_table(
project.adapter.connections.get_bq_table(
relation.database, relation.schema, relation.table
)
)
labels_to_be_created = PROFILE_SCAN_LABELS + list(ORIGINAL_LABELS.keys())
assert set(table.labels.keys()) == set(labels_to_be_created)

0 comments on commit cd6965b

Please sign in to comment.