dbt-data-reliability/integration_tests/tests/test_disable_samples_config.py at 3ba7a44a7fcb5eff48ab3173febe11b671c6fae9 · elementary-data/dbt-data-reliability · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import json

import pytest
from dbt_project import DbtProject

COLUMN_NAME = "sensitive_data"

SAMPLES_QUERY = """
    with latest_elementary_test_result as (
        select id
        from {{ ref("elementary_test_results") }}
        where lower(table_name) = lower('{test_id}')
        order by created_at desc
        limit 1
    )

    select result_row
    from {{ ref("test_result_rows") }}
    where elementary_test_results_id in (select * from latest_elementary_test_result)
"""


@pytest.mark.skip_targets(["clickhouse"])
def test_disable_samples_config_prevents_sampling(
    test_id: str, dbt_project: DbtProject
):
    null_count = 20
    data = [{COLUMN_NAME: None} for _ in range(null_count)]

    columns = [
        {
            "name": COLUMN_NAME,
            "config": {"disable_samples": True},
            "tests": [{"not_null": {}}],
        }
    ]

    test_result = dbt_project.test(
        test_id,
        "not_null",
        columns=columns,
        data=data,
        test_vars={
            "enable_elementary_test_materialization": True,
            "test_sample_row_count": 5,
        },
    )
    assert test_result["status"] == "fail"

    samples = [
        json.loads(row["result_row"])
        for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
    ]
    assert len(samples) == 0


@pytest.mark.skip_targets(["clickhouse"])
def test_disable_samples_false_allows_sampling(test_id: str, dbt_project: DbtProject):
    null_count = 20
    data = [{COLUMN_NAME: None} for _ in range(null_count)]

    columns = [
        {
            "name": COLUMN_NAME,
            "config": {"disable_samples": False},
            "tests": [{"not_null": {}}],
        }
    ]

    test_result = dbt_project.test(
        test_id,
        "not_null",
        columns=columns,
        data=data,
        test_vars={
            "enable_elementary_test_materialization": True,
            "test_sample_row_count": 5,
        },
    )
    assert test_result["status"] == "fail"

    samples = [
        json.loads(row["result_row"])
        for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
    ]
    assert len(samples) == 5
    assert all([row == {COLUMN_NAME: None} for row in samples])


@pytest.mark.skip_targets(["clickhouse"])
def test_disable_samples_config_overrides_pii_tags(
    test_id: str, dbt_project: DbtProject
):
    null_count = 20
    data = [{COLUMN_NAME: None} for _ in range(null_count)]

    columns = [
        {
            "name": COLUMN_NAME,
            "config": {"disable_samples": True, "tags": ["pii"]},
            "tests": [{"not_null": {}}],
        }
    ]

    test_result = dbt_project.test(
        test_id,
        "not_null",
        columns=columns,
        data=data,
        test_vars={
            "enable_elementary_test_materialization": True,
            "test_sample_row_count": 5,
            "disable_samples_on_pii_columns": True,
        },
    )
    assert test_result["status"] == "fail"

    samples = [
        json.loads(row["result_row"])
        for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
    ]
    assert len(samples) == 0