# -*- 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', '')