poyee-data-warehouse/tests/unit/datax/test_mask.py

# -*- coding:utf-8 -*-
import pytest

from dw_base.datax.mask import build_query_sql, resolve_template


def test_static_month_trunc():
    tpl = resolve_template('postgresql', 'month_trunc')
    assert tpl == "TO_CHAR({col}, 'YYYY-MM') AS {col}"


def test_static_md5():
    tpl = resolve_template('postgresql', 'md5')
    assert tpl == 'MD5({col}::text) AS {col}'


def test_dynamic_keep_last_4():
    tpl = resolve_template('postgresql', 'keep_last_4')
    assert tpl == "'****' || RIGHT({col}::text, 4) AS {col}"


def test_dynamic_keep_last_10():
    tpl = resolve_template('postgresql', 'keep_last_10')
    assert 'RIGHT({col}::text, 10)' in tpl


def test_dynamic_keep_first_3():
    tpl = resolve_template('postgresql', 'keep_first_3')
    assert tpl == "LEFT({col}::text, 3) || '****' AS {col}"


def test_build_query_sql_mixed():
    sql = build_query_sql(
        db_type='postgresql',
        columns=['id', 'cert_birthday', 'phone'],
        mask_config={'cert_birthday': 'month_trunc', 'phone': 'keep_last_4'},
        table='public.t',
        where="create_time >= '${start_date}'",
    )
    assert 'SELECT id,' in sql
    assert "TO_CHAR(cert_birthday, 'YYYY-MM') AS cert_birthday" in sql
    assert "RIGHT(phone::text, 4)" in sql
    assert 'FROM public.t' in sql
    assert "WHERE create_time >= '${start_date}'" in sql


def test_invalid_col_name_rejects_injection():
    with pytest.raises(ValueError, match='非法列名'):
        build_query_sql('postgresql', ['id; DROP TABLE x'], {}, 'public.t', '')


def test_unknown_mask_type():
    with pytest.raises(ValueError, match='未知脱敏类型'):
        build_query_sql('postgresql', ['phone'], {'phone': 'bogus_mask'}, 'public.t', '')


def test_mask_col_not_in_columns():
    with pytest.raises(ValueError, match='不在 reader.column 中'):
        build_query_sql('postgresql', ['id'], {'phone': 'md5'}, 'public.t', '')